[Mlir-commits] [mlir] d12d485 - [mlir][sparse] add source materizalization callback for sparse tensor codegen type converter.

[Peiming Liu]

Author: Peiming Liu
Date: 2022-10-20T18:04:57Z
New Revision: d12d4857c5fded57f6333b60cfebb1196b3ac0d3

URL: https://github.com/llvm/llvm-project/commit/d12d4857c5fded57f6333b60cfebb1196b3ac0d3
DIFF: https://github.com/llvm/llvm-project/commit/d12d4857c5fded57f6333b60cfebb1196b3ac0d3.diff

LOG: [mlir][sparse] add source materizalization callback for sparse tensor codegen type converter.

Required by scf.for to achieve 1:N type conversion (See D136314).

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D136367

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index a090e607d5239..e8d0e768852ad 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -745,6 +745,17 @@ class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
 mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
   addConversion(convertSparseTensorType);
+
+  // Required by scf.for 1:N type conversion.
+  addSourceMaterialization([](OpBuilder &builder, RankedTensorType tp,
+                              ValueRange inputs,
+                              Location loc) -> Optional<Value> {
+    if (!getSparseTensorEncoding(tp))
+      // Not a sparse tensor.
+      return llvm::None;
+    // Sparse compiler knows how to cancel out these casts.
+    return genTuple(builder, loc, tp, inputs);
+  });
 }
 
 //===----------------------------------------------------------------------===//