[mlir][sparse] Enable VLA ops in index value generation

[jsetoain]

Summary

Current index value generation uses fixed-length vector ops, this patch
adds an alterantive codegen path compatible with scalable vectors by
using LLVM::StepVectorOp.

[jsetoain]

Clang format

[aartbik]

Thanks for making this work in scalable land!

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

897

Interesting. I don't think we should inject unrealized casts ourselves.
What happens if you simply try to cast this from stepvty to vtp?
Does a later pass do the same then?

[jsetoain]

Replace unrealized cast with index cast

[jsetoain]

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

897

An index cast does actually do the same thing. It was just a debug-induced bad though process on my side :-) Thanks!

[aartbik]

Thanks for keeping the code up-to-date, Javier. Much appreciated!

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

 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.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/SCF.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/Utils/Merger.h"
static Value genIndexValue(Merger &merger, CodeGen &codegen,
   // (1) indices already in vector form, as in ... = ind[lo:hi], good to go, or
   // (2) single index, as in ... = i, must convert to [i, i+1, ...] for inner i.
   unsigned vl = codegen.curVecLength;
   if (vl > 1 && !itype.isa<VectorType>()) {
     Location loc = ival.getLoc();
     VectorType vtp = vectorType(codegen, itype);
     ival = rewriter.create<vector::BroadcastOp>(loc, vtp, ival);
     if (idx == ldx) {
+      Value incr;
+      if (vtp.isScalable()) {
+        Type stepvty = vectorType(codegen, rewriter.getI64Type());
+        Value stepv = rewriter.create<LLVM::StepVectorOp>(loc, stepvty);
+        incr = rewriter.create<arith::IndexCastOp>(loc, vtp, stepv);
+      } else {
         SmallVector<APInt, 4> integers;
         for (unsigned i = 0; i < vl; i++)
           integers.push_back(APInt(/*width=*/64, i));
         auto values = DenseElementsAttr::get(vtp, integers);
-      Value incr = rewriter.create<arith::ConstantOp>(loc, vtp, values);
+        incr = rewriter.create<arith::ConstantOp>(loc, vtp, values);
+      }
       ival = rewriter.create<arith::AddIOp>(loc, ival, incr);
     }
   }
   return ival;
 }
 
 /// Recursively generates tensor expression.
 static Value genExp(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,