[mlir][sparse] Add codegen for expand op.

[aartbik]

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

467

keep the original comment

// Determine the size for access expansion.

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

181

rewriting rules -> conversion rules (or codegen rules)

184

since you mention dialect at L180, L184 can now be removed, just as above

mlir/test/Dialect/SparseTensor/codegen.mlir

317

how about a single letter name too, so it aligns better

319

identation?

323

maybe make this 8x10 so we see what size gets used!

328

let's also add a test with ?x? dimensions and parameters in the bufferization.alloc_tensor(...) so we get a dimOp operator. Perhaps with a column-wise access, so we see that it swaps?

[aartbik]

mlir/test/Dialect/SparseTensor/codegen.mlir

323

and also a CSC with constants

[Peiming]

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

471

This looks very similar to the original implementation, anyway that we can merge these two?

[aartbik]

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

467

Since DimOp operates on the original dimension, and you want the innermost stored dimension, we need to apply toOrig() to the computed rank-1 here. This will be an example how much better "codegen" composes, since we will fold this much more often than we did in the "conversion" case.

[Peiming]

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

467

Just out of curiosity: Is it okay to create DimOp here? since DimOp is illegal and will need to be converted in the same pass as well? Will MLIR recursively legalizes DimOp? If so, why can't we do that in the conversion pass?

Also, will it have the same bug as https://reviews.llvm.org/D133472, you might want a toOrig conversion here as well?

[Peiming]

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

467

Uh oh, seems my comment was just repeating it.

[bixia]

Address review comments.

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

467

Put the comment back and after fix the bug that is fixed in D133472

mlir/test/Dialect/SparseTensor/codegen.mlir

323

I added CSC with dynamic value, combining this with the request below. PTAL.

[bixia]

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

467

Yes, DimOp will be legalized.

471

Good question, I was going to ask whether we care about code sharing between conversion and codegen paths. But will refactor this after we have a decision.

[bixia]

Modified the comment to align with D133512.

[bixia]

Remove a temp variable.

[aartbik]

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

435

this comment is not really valid, since you are not using the utility

please move over the comment in the revision "conversion" case, stating we rewrite recursively on the *original* tensor (hence op.getTensor and not adaptor.getTensor)

mlir/test/Dialect/SparseTensor/codegen.mlir

333

do we have a pass that does the store/load forwarding ;-)

[aartbik]

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

435

scratch this, you already did that

[bixia]

Update the tests to align with D133512.

[bixia]

mlir/test/Dialect/SparseTensor/codegen.mlir

333

Yes. Another reason for switching away from the supporting lib.

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

let description = [{
     ```mlir
     TBD
     ```
   }];
   let constructor = "mlir::createSparseTensorCodegenPass()";
   let dependentDialects = [
     "arith::ArithmeticDialect",
     "bufferization::BufferizationDialect",
+    "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
 }
 
 def SparseTensorStorageExpansion : Pass<"sparse-tensor-storage-expansion", "ModuleOp"> {
   let summary = "Expand compounded sparse tensor storage into individual SSA values";

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

 // compiler optimization of the generated code.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
     Location loc = op->getLoc();
     Value tuple = adaptor.getTensor();
     unsigned i = tuple.getType().cast<TupleType>().size() - 1; // last
     rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
     return success();
   }
 };
 
 /// Sparse codegen rule for tensor rematerialization.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (op.getHasInserts()) {
       // Finalize any pending insertions.
       // TODO: implement
     }
     rewriter.replaceOp(op, adaptor.getOperands());
     return success();
   }
 };
 
+
+/// Sparse codegen rule for the expand op.
+class SparseExpandConverter : public OpConversionPattern<ExpandOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+    ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
+    Type eltType = srcType.getElementType();
+    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());
+    Value src = op.getTensor();
+    Value sz = rewriter.create<tensor::DimOp>(loc, src, srcType.getRank() - 1);
+    // Generate a memref for `sz` elements of type `t`.
+    auto genAlloc = [&](Type t) {
+      auto memTp = MemRefType::get({ShapedType::kDynamicSize}, t);
+      return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{sz});
+    };
+    // Allocate temporary buffers for values, filled-switch, and indices.
+    // We do not use stack buffers for this, since the expanded size may
+    // be rather large (as it envelops a single expanded dense dimension).
+    Value values = genAlloc(eltType);
+    Value filled = genAlloc(boolType);
+    Value indices = genAlloc(idxType);
+    Value zero = constantZero(rewriter, loc, idxType);
+    // Reset the values/filled-switch to all-zero/false. Note that this
+    // introduces an O(N) operation into the computation, but this reset
+    // operation is amortized over the innermost loops for the access
+    // pattern expansion. As noted in the operation doc, we would like
+    // to amortize this setup cost even between kernels.
+    rewriter.create<linalg::FillOp>(
+        loc, ValueRange{constantZero(rewriter, loc, eltType)},
+        ValueRange{values});
+    rewriter.create<linalg::FillOp>(
+        loc, ValueRange{constantZero(rewriter, loc, boolType)},
+        ValueRange{filled});
+    // Replace expansion op with these buffers and initial index.
+    assert(op.getNumResults() == 4);
+    rewriter.replaceOp(op, {values, filled, indices, zero});
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into an actual buffer.
 //===----------------------------------------------------------------------===//
 
 mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
   addConversion(convertSparseTensorType);
 }
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseDimOpConverter, SparseCastConverter,
-               SparseTensorAllocConverter, SparseTensorDeallocConverter,
-               SparseToPointersConverter, SparseToIndicesConverter,
-               SparseToValuesConverter, SparseTensorLoadConverter>(
-      typeConverter, patterns.getContext());
+               SparseExpandConverter, SparseTensorAllocConverter,
+               SparseTensorDeallocConverter, SparseToPointersConverter,
+               SparseToIndicesConverter, SparseToValuesConverter,
+               SparseTensorLoadConverter>(typeConverter, patterns.getContext());
 }

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

void runOnOperation() override {
     target.addDynamicallyLegalOp<bufferization::AllocTensorOp>(
         [&](bufferization::AllocTensorOp op) {
           return converter.isLegal(op.getType());
         });
     target.addDynamicallyLegalOp<bufferization::DeallocTensorOp>(
         [&](bufferization::DeallocTensorOp op) {
           return converter.isLegal(op.getTensor().getType());
         });
+    // The following operations and dialects may be introduced by the
+    // rewriting rules, and are therefore marked as legal.
+    target.addLegalOp<linalg::FillOp>();
+
     // Legal dialects may occur in generated code.
     target.addLegalDialect<arith::ArithmeticDialect,
                            bufferization::BufferizationDialect,
                            memref::MemRefDialect, scf::SCFDialect>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     populateCallOpTypeConversionPattern(patterns, converter);
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,

mlir/test/Dialect/SparseTensor/codegen.mlir

 //      CHECK-CODEGEN: %[[T:.*]] = sparse_tensor.storage(%[[A0]], %[[A1]])
 //      CHECK-CODEGEN: return %[[T]] : tuple<memref<3xindex>, memref<?xf64>>
 //      CHECK-STORAGE: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf64>
 func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {
   %0 = bufferization.alloc_tensor() : tensor<10x20x30xf64, #Dense3D>
   %1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D>
   return %1 : tensor<10x20x30xf64, #Dense3D>
 }
+
+// CHECK-LABEL:   func.func @sparse_expansion()
+// CHECK:           %[[A:.*]] = memref.alloc() : memref<8xf64>
+// CHECK:           %[[B:.*]] = memref.alloc() : memref<8xi1>
+// CHECK:           %[[CStatic:.*]] = memref.alloc() : memref<8xindex>
+// CHECK:           %[[C:.*]] = memref.cast %[[CStatic]] : memref<8xindex> to memref<?xindex>
+// CHECK-DAG:           linalg.fill ins(%{{.*}}  : f64) outs(%[[A]] : memref<8xf64>)
+// CHECK-DAG:           linalg.fill ins(%{{.*}}  : i1) outs(%[[B]] : memref<8xi1>)
+// CHECK:           return %[[C]] : memref<?xindex>
+func.func @sparse_expansion() -> memref<?xindex> {
+  %0 = bufferization.alloc_tensor() : tensor<8x8xf64, #CSR>
+  %values, %filled, %added, %count = sparse_tensor.expand %0
+    : tensor<8x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return %added : memref<?xindex>
+}
+