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

[aartbik]

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

585

keep the original comment

// Determine the size for access expansion.

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

178–180

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

179

rewriting rules -> conversion rules (or codegen rules)

mlir/test/Dialect/SparseTensor/codegen.mlir

293

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

295

identation?

299

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

304

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

299

and also a CSC with constants

[Peiming]

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

589

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

[aartbik]

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

585

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

585

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

585

Uh oh, seems my comment was just repeating it.

[bixia]

Address review comments.

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

585

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

mlir/test/Dialect/SparseTensor/codegen.mlir

299

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

[bixia]

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

585

Yes, DimOp will be legalized.

589

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

494

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

309

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

[aartbik]

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

494

scratch this, you already did that

[bixia]

Update the tests to align with D133512.

[bixia]

mlir/test/Dialect/SparseTensor/codegen.mlir

309

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",
   ];
 }
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

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(SourceOp op, OpAdaptor adaptor,
     unsigned idx = Base::getIndexForOp(tuple, op);
     auto fields = tuple.getInputs();
     assert(idx < fields.size());
     rewriter.replaceOp(op, fields[idx]);
     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());
+    // Determine the size for access expansion (always the innermost stored
+    // dimension size, translated back to original dimension). Note that we
+    // recursively rewrite the new DimOp on the **original** tensor.
+    auto enc = getSparseTensorEncoding(srcType);
+    unsigned innerDim = srcType.getRank() - 1;
+    if (AffineMap p = enc.getDimOrdering())
+      innerDim = p.getDimPosition(innerDim);
+    Value sz = rewriter.create<tensor::DimOp>(loc, op.getTensor(), innerDim);
+    // 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();
+  }
+};
+
 /// Sparse codegen rule for pointer accesses.
 class SparseToPointersConverter
     : public SparseGetterOpConverter<ToPointersOp, SparseToPointersConverter> {
 public:
   using SparseGetterOpConverter::SparseGetterOpConverter;
   // Callback for SparseGetterOpConverter.
   static unsigned getIndexForOp(UnrealizedConversionCastOp /*tuple*/,
                                 ToPointersOp op) {
mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   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, SparseCallConverter, SparseDimOpConverter,
-               SparseCastConverter, SparseTensorAllocConverter,
-               SparseTensorDeallocConverter, SparseToPointersConverter,
-               SparseToIndicesConverter, SparseToValuesConverter,
-               SparseTensorLoadConverter>(typeConverter, patterns.getContext());
+               SparseCastConverter, 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());
         });
-    // Legal dialects may occur in generated code.
+    // The following operations and dialects may be introduced by the
+    // codegen rules, and are therefore marked as legal.
+    target.addLegalOp<linalg::FillOp>();
     target.addLegalDialect<arith::ArithmeticDialect,
                            bufferization::BufferizationDialect,
                            memref::MemRefDialect, scf::SCFDialect>();
     target.addLegalOp<UnrealizedConversionCastOp>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,

mlir/test/Dialect/SparseTensor/codegen.mlir

 //              CHECK: %[[A:.*]] = memref.alloc() : memref<6000xf64>
 //              CHECK: %[[A1:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>
 //              CHECK: 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_expansion1()
+//         CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64>
+//         CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1>
+//         CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex>
+//         CHECK: %[[D:.*]] = memref.cast %[[C]] : 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 %[[D]] : memref<?xindex>
+func.func @sparse_expansion1() -> memref<?xindex> {
+  %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSR>
+  %values, %filled, %added, %count = sparse_tensor.expand %0
+    : tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return %added : memref<?xindex>
+}
+
+//   CHECK-LABEL: func.func @sparse_expansion2()
+//         CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64>
+//         CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1>
+//         CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex>
+//         CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex>
+//     CHECK-DAG: linalg.fill ins(%{{.*}}  : f64) outs(%[[A]] : memref<4xf64>)
+//     CHECK-DAG: linalg.fill ins(%{{.*}}  : i1) outs(%[[B]] : memref<4xi1>)
+//         CHECK: return %[[D]] : memref<?xindex>
+func.func @sparse_expansion2() -> memref<?xindex> {
+  %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSC>
+  %values, %filled, %added, %count = sparse_tensor.expand %0
+    : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return %added : memref<?xindex>
+}
+
+//   CHECK-LABEL: func.func @sparse_expansion3(
+//    CHECK-SAME: %[[D0:.*]]: index,
+//    CHECK-SAME: %{{.*}}: index) -> memref<?xindex> {
+//         CHECK: %[[C1:.*]] = arith.constant 1 : index
+//         CHECK: %[[S0:.*]] = memref.alloc() : memref<2xindex>
+//         CHECK: memref.store %[[D0]], %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
+//         CHECK: %[[D1:.*]] = memref.load %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
+//         CHECK: %[[V:.*]] = memref.alloc(%[[D1]]) : memref<?xf64>
+//         CHECK: %[[B:.*]] = memref.alloc(%[[D1]]) : memref<?xi1>
+//         CHECK: %[[D:.*]] = memref.alloc(%[[D1]]) : memref<?xindex>
+//         CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[V]] : memref<?xf64>)
+//         CHECK: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
+//         CHECK: return %[[D]] : memref<?xindex>
+func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> {
+  %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC>
+  %values, %filled, %added, %count = sparse_tensor.expand %0
+    : tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return %added : memref<?xindex>
+}