[Mlir-commits] [mlir] [mlir][sparse_tensor] Handle block argument for sparse_tensor.expand codegen (PR #205487)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: jpwang (jjppp)

<details>
<summary>Changes</summary>

Handle the case where the source tensor of `sparse_tensor.expand` is a block argument in SparseExpandConverter. Previously it unconditionally uses `getDefiningOp()`, which is `null` for block arguments and causes a crash.

A regression test is added to cover `sparse_tensor.expand` on a function argument.

Fixes #<!-- -->204712 

---
Full diff: https://github.com/llvm/llvm-project/pull/205487.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp (+5-1) 
- (added) mlir/test/Dialect/SparseTensor/sparse_tensor_codegen_expand.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 2a57fcd0b36a8..38558777d537c 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -912,7 +912,11 @@ class SparseExpandConverter : public OpConversionPattern<ExpandOp> {
     Type boolType = rewriter.getIntegerType(1);
     Type idxType = rewriter.getIndexType();
     // All initialization should be done on entry of the loop nest.
-    rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
+    if (isa<BlockArgument>(op.getTensor())) {
+      rewriter.setInsertionPointToStart(op->getBlock());
+    } else {
+      rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
+    }
 
     // Determine the size for access expansion (always the innermost stored
     // level size).
diff --git a/mlir/test/Dialect/SparseTensor/sparse_tensor_codegen_expand.mlir b/mlir/test/Dialect/SparseTensor/sparse_tensor_codegen_expand.mlir
new file mode 100644
index 0000000000000..15cb2599ccda6
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_tensor_codegen_expand.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -sparse-tensor-codegen | FileCheck %s
+
+#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed, d1 : compressed) }>
+
+// This test verifies that sparse_tensor.expand codegen handles the case where
+// the input sparse tensor is a function block argument.
+// CHECK-LABEL: func.func @sparse_expansion(
+// CHECK-SAME: %[[A0:.*]]: memref<?xindex>, %[[A1:.*]]: memref<?xindex>, %[[A2:.*]]: memref<?xindex>, %[[A3:.*]]: memref<?xindex>, %[[A4:.*]]: memref<?xf64>, %[[A5:.*]]: !sparse_tensor.storage_specifier<#sparse>) -> index
+// CHECK-NOT: sparse_tensor.expand
+
+module {
+  func.func @sparse_expansion(%arg0: tensor<8x8xf64, #sparse>) -> index {
+    %values, %filled, %added, %count = sparse_tensor.expand %arg0
+      : tensor<8x8xf64, #sparse> to memref<?xf64>, memref<?xi1>, memref<?xindex>
+    return %count : index
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/205487