[MLIR][SparseTensor] Reject unsupported symbolic demapping

mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp

@@ -38,6 +38,12 @@ struct DemapInsRewriter : public OpRewritePattern<SourceOp> {
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
 
+    for (Value in : op->getOperands())
+      if (auto stt = tryGetSparseTensorType(in);
+          stt && !stt->isIdentity() &&
+          stt->getEncoding().getDimToLvl().getNumSymbols() != 0)
+        return failure();
+
     // Demaps non-trivial inputs.
     bool changed = false;
     SmallVector<Value> deMappedIns(op->getOperands());
@@ -605,6 +611,8 @@ struct TensorAllocDemapper : public OpRewritePattern<AllocOp> {
 
     Location loc = op.getLoc();
     auto stt = getSparseTensorType(op.getResult());
+    if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
+      return failure();
 
     SmallVector<Value> maxDimCrds;
     maxDimCrds.reserve(stt.getDimRank());
@@ -674,6 +682,8 @@ struct SparseAssembleDemapper : public OpRewritePattern<AssembleOp> {
 
     assert(hasAnySparseResult(op));
     auto stt = getSparseTensorType(op.getResult());
+    if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
+      return failure();
     rewriter.modifyOpInPlace(
         op, [&op, &stt]() { op.getResult().setType(stt.getDemappedType()); });
     rewriter.setInsertionPointAfter(op);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

@@ -1340,6 +1340,8 @@ struct CrdTranslateRewriter : public OpRewritePattern<CrdTranslateOp> {
     AffineMap map = op.getDirection() == CrdTransDirectionKind::dim2lvl
                         ? op.getEncoder().getDimToLvl()
                         : op.getEncoder().getLvlToDim();
+    if (!map || map.getNumSymbols() != 0)
+      return failure();
 
     SmallVector<Value> outCrds;
     for (AffineExpr result : map.getResults()) {

mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir

@@ -1,8 +1,7 @@
 // RUN: mlir-opt %s -split-input-file -sparsification-and-bufferization -verify-diagnostics | FileCheck %s
 
-// XFAIL: mlir-expensive-checks
-
-// Tests that mlir-opt does not crash when parsing sparse tensor encodings with symbols.
+// Tests that processing sparse tensor encodings with symbols does not crash and
+// reports a diagnostic when lowering is unsupported.
 
 // CHECK-DAG: #[[$SPARSE_0:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d1 : dense, d2 : compressed) }>
 // CHECK-DAG: #[[$SPARSE_1:.*]] = #sparse_tensor.encoding<{ map = [s0](d0, d1) -> (d0 * (s0 * 3) : dense, d0 : dense, d1 : compressed) }>
@@ -47,7 +46,7 @@ func.func @tensor_convert() -> memref<?xindex> {
     tensor.yield %val : f32
   } : tensor<32x32xf32>
 
-  // expected-error@+1 {{Level size mismatch between source/dest tensors}}
+  // expected-error@+1 {{failed to legalize operation 'bufferization.alloc_tensor'}}
   %J = sparse_tensor.convert %I : tensor<32x32xf32> to tensor<32x32xf32, #Sparse>
 
   %result = sparse_tensor.positions %J { level = 0 : index }

mlir/test/Dialect/SparseTensor/sparse_foreach.mlir

@@ -170,3 +170,16 @@ func.func @foreach_bcoo(%A: tensor<4x4x4xf64, #BCOO>) {
   }
   return
 }
+
+#Symbolic = #sparse_tensor.encoding<{
+  map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @symbolic_crd_translate(
+// CHECK: sparse_tensor.crd_translate dim_to_lvl
+func.func @symbolic_crd_translate(%i: index, %j: index)
+    -> (index, index, index) {
+  %l0, %l1, %l2 = sparse_tensor.crd_translate dim_to_lvl [%i, %j]
+      as #Symbolic : index, index, index
+  return %l0, %l1, %l2 : index, index, index
+}

mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir

@@ -129,3 +129,21 @@ func.func @sparse_disassemble_reinterpret_map(%sp : tensor<2x4xf64, #BSR>,
                                  -> (tensor<?xindex>, tensor<?xindex>), tensor<?xf64>, (index, index), index
   return %rd, %rp, %ri : tensor<?xf64>, tensor<?xindex>, tensor<?xindex>
 }
+
+// -----
+
+#Symbolic = #sparse_tensor.encoding<{
+  map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @sparse_assemble_symbolic_map(
+// CHECK: %[[ASSEMBLED:.*]] = sparse_tensor.assemble {{.*}} to tensor<32x32xf32, #sparse{{[0-9]*}}>
+// CHECK-NEXT: return %[[ASSEMBLED]] : tensor<32x32xf32, #sparse{{[0-9]*}}>
+func.func @sparse_assemble_symbolic_map(
+    %val: tensor<?xf32>, %pos: tensor<?xindex>, %crd: tensor<?xindex>)
+    -> tensor<32x32xf32, #Symbolic> {
+  %0 = sparse_tensor.assemble (%pos, %crd), %val
+      : (tensor<?xindex>, tensor<?xindex>), tensor<?xf32>
+        to tensor<32x32xf32, #Symbolic>
+  return %0 : tensor<32x32xf32, #Symbolic>
+}