[mlir][sparse] code cleanup using the assumption that dim2lvl maps ar… (#72894)

…e simplified.

Author: PeimingLiu

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

@@ -745,8 +745,8 @@ class SparseTensorAllocConverter
     const auto resType = getSparseTensorType(op);
     if (!resType.hasEncoding())
       return failure();
-    Location loc = op.getLoc();
 
+    Location loc = op.getLoc();
     // Deal with copy.
     if (op.getCopy()) {
       auto desc = getDescriptorFromTensorTuple(adaptor.getCopy());
@@ -768,16 +768,14 @@ class SparseTensorAllocConverter
       return success();
     }
 
-    // Construct the dim/lvl sizes and the (unused) dim2lvl/lvl2dim buffers.
-    SmallVector<Value> dimSizesValues;
+    if (!resType.isIdentity()) {
+      return rewriter.notifyMatchFailure(
+          op, "try run --sparse-reinterpret-map before codegen");
+    }
+    // Level size equals to dimension size since lvl2dim map is an identity map.
     SmallVector<Value> lvlSizesValues;
-    Value dimSizesBuffer;
-    Value dim2lvlBuffer;
-    Value lvl2dimBuffer;
     createDimSizes(rewriter, loc, resType, adaptor.getDynamicSizes(),
-                   dimSizesValues);
-    genMapBuffers(rewriter, loc, resType, dimSizesValues, dimSizesBuffer,
-                  lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer);
+                   /*dimSizesValues=*/lvlSizesValues);
 
     // Construct allocation for each field.
     Value sizeHint = op.getSizeHint();
@@ -809,19 +807,17 @@ class SparseTensorEmptyConverter : public OpConversionPattern<tensor::EmptyOp> {
     const auto resType = getSparseTensorType(op);
     if (!resType.hasEncoding())
       return failure();
-    Location loc = op.getLoc();
 
-    // Construct the dim/lvl sizes and the (unused) dim2lvl/lvl2dim buffers.
-    SmallVector<Value> dimSizesValues;
+    if (!resType.isIdentity()) {
+      return rewriter.notifyMatchFailure(
+          op, "try run --sparse-reinterpret-map before codegen");
+    }
+
+    Location loc = op.getLoc();
+    // Level size equals to dimension size since lvl2dim map is an identity map.
     SmallVector<Value> lvlSizesValues;
-    Value dimSizesBuffer;
-    Value dim2lvlBuffer;
-    Value lvl2dimBuffer;
     createDimSizes(rewriter, loc, resType, adaptor.getDynamicSizes(),
-                   dimSizesValues);
-    genMapBuffers(rewriter, loc, resType, dimSizesValues, dimSizesBuffer,
-                  lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer);
-
+                   /*dimSizesValues=*/lvlSizesValues);
     // Construct allocation for each field.
     Value sizeHint; // none
     SmallVector<Value> fields;

mlir/test/Dialect/SparseTensor/codegen.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --lower-sparse-ops-to-foreach --lower-sparse-foreach-to-scf --sparse-tensor-codegen  --canonicalize -cse | FileCheck %s
+// RUN: mlir-opt %s --lower-sparse-ops-to-foreach --lower-sparse-foreach-to-scf --sparse-reinterpret-map --sparse-tensor-codegen  --canonicalize -cse | FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>