[mlir][sparse][simd] only accept proper unit stride subscripts

[bixia]

mlir/test/Dialect/SparseTensor/sparse_vector_concat.mlir

1

no needed.

[aartbik]

removed debug/print

[aartbik]

rebased with main (including latest vector reduction test)

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

 /// For the scalar case, subscripts simply zero extend narrower indices
 /// into 64-bit values before casting to an index type without a performance
 /// penalty. Indices that already are 64-bit, in theory, cannot express the
 /// full range since the LLVM backend defines addressing in terms of an
 /// unsigned pointer/signed index pair.
 static bool vectorizeSubscripts(PatternRewriter &rewriter, scf::ForOp forOp,
                                 VL vl, ValueRange subs, bool codegen,
                                 Value vmask, SmallVectorImpl<Value> &idxs) {
+  unsigned d = 0;
+  unsigned dim = subs.size();
   for (auto sub : subs) {
-    // Invariant/loop indices simply pass through.
-    if (sub.dyn_cast<BlockArgument>() ||
+    bool innermost = ++d == dim;
+    // Invariant subscripts in outer dimensions simply pass through.
+    // Note that we rely on LICM to hoist loads where all subscripts
+    // are invariant in the innermost loop.
+    if (sub.getDefiningOp() &&
         sub.getDefiningOp()->getBlock() != &forOp.getRegion().front()) {
+      if (innermost)
+        return false;
+      if (codegen)
+        idxs.push_back(sub);
+      continue; // success so far
+    }
+    // Invariant block arguments (including outer loop indices) in outer
+    // dimensions simply pass through. Direct loop indices in the
+    // innermost loop simply pass through as well.
+    if (auto barg = sub.dyn_cast<BlockArgument>()) {
+      bool invariant = barg.getOwner() != &forOp.getRegion().front();
+      if (invariant == innermost)
+        return false;
       if (codegen)
         idxs.push_back(sub);
       continue; // success so far
     }
     // Look under the hood of casting.
     auto cast = sub;
     while (true) {
       if (auto icast = cast.getDefiningOp<arith::IndexCastOp>())
for (auto sub : subs) {
     // zero extending the elements into 64-bit loses some performance since
     // the 32-bit indexed gather/scatter is more efficient than the 64-bit
     // index variant (if the negative 32-bit index space is unused, the
     // enableSIMDIndex32 flag can preserve this performance). For 64-bit
     // values, there is no good way to state that the indices are unsigned,
     // which creates the potential of incorrect address calculations in the
     // unlikely case we need such extremely large offsets.
     if (auto load = cast.getDefiningOp<memref::LoadOp>()) {
+      if (!innermost)
+        return false;
       if (codegen) {
         SmallVector<Value> idxs2(load.getIndices()); // no need to analyze
         Location loc = forOp.getLoc();
         Value vload =
             genVectorLoad(rewriter, loc, vl, load.getMemRef(), idxs2, vmask);
         Type etp = vload.getType().cast<VectorType>().getElementType();
         if (!etp.isa<IndexType>()) {
           if (etp.getIntOrFloatBitWidth() < 32)
             vload = rewriter.create<arith::ExtUIOp>(
                 loc, vectorType(vl, rewriter.getI32Type()), vload);
           else if (etp.getIntOrFloatBitWidth() < 64 && !vl.enableSIMDIndex32)
             vload = rewriter.create<arith::ExtUIOp>(
                 loc, vectorType(vl, rewriter.getI64Type()), vload);
         }
         idxs.push_back(vload);
       }
       continue; // success so far
     }
     // Address calculation 'i = add inv, idx' (after LICM).
     if (auto load = cast.getDefiningOp<arith::AddIOp>()) {
       Value inv = load.getOperand(0);
       Value idx = load.getOperand(1);
-      if (!inv.dyn_cast<BlockArgument>() &&
+      if (inv.getDefiningOp() &&
           inv.getDefiningOp()->getBlock() != &forOp.getRegion().front() &&
           idx.dyn_cast<BlockArgument>()) {
+        if (!innermost)
+          return false;
         if (codegen)
           idxs.push_back(
               rewriter.create<arith::AddIOp>(forOp.getLoc(), inv, idx));
         continue; // success so far
       }
     }
     return false;
   }

mlir/test/Dialect/SparseTensor/sparse_vector_concat.mlir

+// RUN: mlir-opt %s --sparse-compiler="enable-runtime-library=false vl=2 reassociate-fp-reductions=true enable-index-optimizations=true" --print-before-all --debug
+
+#MAT_D_C = #sparse_tensor.encoding<{
+  dimLevelType = ["dense", "compressed"]
+}>
+
+#MAT_C_C_P = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+#MAT_C_D_P = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "dense" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+//
+// Ensures only last loop is vectorized
+// (vectorizing the others would crash).
+//
+// CHECK-LABEL: llvm.func @foo
+// CHECK:       llvm.intr.masked.load
+// CHECK:       llvm.intr.masked.scatter
+//
+func.func @foo(%arg0: tensor<2x4xf64, #MAT_C_C_P>,
+               %arg1: tensor<3x4xf64, #MAT_C_D_P>,
+           %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> {
+  %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
+       : tensor<2x4xf64, #MAT_C_C_P>, tensor<3x4xf64, #MAT_C_D_P>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64>
+  return %0 : tensor<9x4xf64>
+}