[mlir][sparse] Fix rewriting for convert op and concatenate op.

[bixia]

Summary

Fix a problem in convert op rewriting where it used the original index for
ToIndicesOp.

Extend the concatenate op rewriting to handle dense destination and dynamic
shape destination.

Make the concatenate op integration test run on the codegen path.

[Peiming]

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

438

I think concatenate verifier ensures that only concatenate dimension can be dynamic.

So probably you do not need to iterate over shape here.

[Peiming]

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

438

Oh, I am wrong. all the input are statically sized, only output tensor can be dynamically sized...

[Peiming]

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

435

nit: I learned this from Alex, who suggest me not specify the number of stack elements in smallvector (unless you have strong reasons for it), maybe we should start to follow the new practice.

[Peiming]

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

441

Maybe you should call createOrFoldDimOp instead?

[bixia]

Use createOrFoldDimOp and remove N from SmallVector<>.

[aartbik]

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

435

Yeah, in the past we used 4 for things that typically hold a rank, but I am perfectly okay letting the CalculateSmallVectorDefaultInlinedElements() do its magic.

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Support/LLVM.h"
if (rank > 1) {
     // level should the COO has a unique level at the end. Ends by a unordered
     // unique singleton level unless the tensor rank is 1.
     dims.push_back(DimLevelType::SingletonNo);
   }
   SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(src);
   // TODO: Maybe pick the bitwidth based on input/output tensors (probably the
   // largest one among them) in the original operation instead of using the
   // default value.
+  unsigned pointerBitWidth = encSrc ? encSrc.getPointerBitWidth() : 0;
+  unsigned indexBitWidth = encSrc ? encSrc.getIndexBitWidth() : 0;
   auto enc = SparseTensorEncodingAttr::get(
       ctx, dims, AffineMap::getMultiDimIdentityMap(rank, ctx), AffineMap(),
-      encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
+      pointerBitWidth, indexBitWidth);
   return RankedTensorType::get(src.getShape(), src.getElementType(), enc);
 }
 
 /// Collects the dynamic dimension sizes for `tp` with the assumption that
 /// `sizes` are the dimension sizes for the type. Stores the dynamic dimension
 /// sizes to dynSizes.
 static void getDynamicSizes(RankedTensorType tp,
                             const SmallVectorImpl<Value> &sizes,

 };
 
 struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(ConcatenateOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     auto rtp = op.getType().cast<RankedTensorType>();
-    // TODO: Build the output shape if needed.
-    assert(rtp.hasStaticShape());
-    auto rank = rtp.getRank();
     size_t conDim = op.getDimension().getZExtValue();
+    SmallVector<Value> dynSizes;
+    if (!rtp.hasStaticShape()) {
+      ArrayRef<int64_t> rShape = rtp.getShape();
+      for (const auto &d : llvm::enumerate(rShape)) {
+        if (d.value() == ShapedType::kDynamicSize) {
+          Value v =
+              createOrFoldDimOp(rewriter, loc, op.getOperand(0), d.index());
+          rewriter.create<tensor::DimOp>(loc, op.getOperand(0), d.index());
+          for (const auto &opnd : op.getOperands().drop_front()) {
+            Value t = createOrFoldDimOp(rewriter, loc, opnd, d.index());
+            v = rewriter.create<arith::AddIOp>(loc, v, t);
+          }
+          dynSizes.push_back(v);
+        }
+      }
+    }
+
     // %t = concatenate %s1, %s2, %s3 {dim = 1}
     // ==>
     // %tmp = bufferization.alloc_tensor : unordered COO
     // foreach in %s1 : insert d0, d1, %tmp
     // foreach in %s2 : insert d0, d1 + size(s1), %tmp
     // foreach in %s3 : insert d0, d1 + size(s1) + size(s2), %tmp
     // %t = sparse_tensor.cast %tmp
     auto cooTp = getUnorderedCOOFromType(rtp);
     auto cooBuffer =
-        rewriter.create<AllocTensorOp>(loc, cooTp, ValueRange()).getResult();
-
+        rewriter.create<AllocTensorOp>(loc, cooTp, dynSizes).getResult();
+    auto rank = rtp.getRank();
     Value offset = constantIndex(rewriter, loc, 0);
     ForeachOp foreachOp;
     for (Value input : op.getInputs()) {
-      // Builds the indexing map.
-
       // Build a for op for each input tensor to append new values into the
       // output tensor.
       foreachOp = rewriter.create<ForeachOp>(
           loc, input, cooBuffer,
           [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
               ValueRange reduc) {
             SmallVector<Value, 4> indices;
             for (int64_t i = 0; i < rank; i++) {
               Value idx = args[i];
               if (i == static_cast<int64_t>(conDim))
                 // transform coordinates on matching dim
                 idx = builder.create<arith::AddIOp>(loc, idx, offset);
               indices.push_back(idx);
             }
-            auto t = builder.create<InsertOp>(loc, v, reduc.front(), indices);
-            builder.create<sparse_tensor::YieldOp>(loc, t);
+            Value cond = genIsNonzero(rewriter, loc, v);
+            scf::IfOp ifOp = builder.create<scf::IfOp>(
+                loc, TypeRange(reduc.front().getType()), cond, /*else*/ true);
+            builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+            Value t = builder.create<InsertOp>(loc, v, reduc.front(), indices);
+            rewriter.create<scf::YieldOp>(loc, t);
+            rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
+            rewriter.create<scf::YieldOp>(loc, reduc.front());
+            rewriter.setInsertionPointAfter(ifOp);
+            rewriter.create<sparse_tensor::YieldOp>(loc, ifOp.getResult(0));
           });
       // Accumulates the offset. Note that only static-shaped inputs are allowed
       // by concatenate op verifier, which saves us from computing the offset
       // dynamically.
       auto d = input.getType().cast<RankedTensorType>().getShape()[conDim];
       assert(!ShapedType::isDynamic(d));
       offset = rewriter.create<arith::AddIOp>(loc, offset,
                                               constantIndex(rewriter, loc, d));
LogicalResult sparse2SparseRewrite(ConvertOp op,
     auto indTp =
         MemRefType::get(dynShape, getIndexOverheadType(rewriter, encSrc));
     uint64_t rank = dstTp.getRank();
     // Gather the indices-arrays in the dst tensor storage order.
     SmallVector<Value, 4> xs(rank, Value());
     for (uint64_t i = 0; i < rank; i++) {
       uint64_t orgDim = toOrigDim(encSrc, i);
       xs[toStoredDim(encDst, orgDim)] = rewriter.create<ToIndicesOp>(
-          loc, indTp, src, rewriter.getIndexAttr(orgDim));
+          loc, indTp, src, rewriter.getIndexAttr(i));
     }
 
     // Retrieve NNZ.
     Value nnz = rewriter.create<NumberOfEntriesOp>(loc, src);
     nnz =
         rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), nnz);
 
     // Retrieve the values-array.

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate.mlir

-// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=true | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+//
+// Do the same run, but now with direct IR generation.
+// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=false | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 #MAT_C_C = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 #MAT_D_C = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
 #MAT_C_D = #sparse_tensor.encoding<{dimLevelType = ["compressed", "dense"]}>