[Mlir-commits] [mlir] [mlir][affine] Support vectorization with the step size exceeding 2^32-1 (PR #66026)

[llvmlistbot at llvm.org]

llvmbot wrote:

@llvm/pr-subscribers-mlir

<details>
<summary>Changes</summary>

Since the step size in `affine.for` is defined as `int64_t`, we should be able to keep the precision as it is with using the `int64_t` type int the vectorization pass. 
--
Full diff: https://github.com/llvm/llvm-project/pull/66026.diff

2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+1-1) 
- (modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir (+13) 


<pre>
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 072e858220feae3..222536756f9fa75 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1291,7 +1291,7 @@ static Operation *vectorizeAffineForOp(AffineForOp forOp,
   // If we are vectorizing a vector dimension, compute a new step for the new
   // vectorized loop using the vectorization factor for the vector dimension.
   // Otherwise, propagate the step of the scalar loop.
-  unsigned newStep;
+  int64_t newStep;
   if (isLoopVecDim) {
     unsigned vectorDim = loopToVecDimIt->second;
     assert(vectorDim < strategy.vectorSizes.size() && "vector dim overflow");
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
index 9244604128cb723..f4f056745c00e98 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
@@ -684,3 +684,16 @@ func.func @vec_vecdim_reduction_rejected(%in: memref<256x512xf32>, %out: memref<
 
 // CHECK-LABEL: @vec_vecdim_reduction_rejected
 // CHECK-NOT: vector
+
+// -----
+
+// CHECK-LABEL: @large_step_size
+// Support the step size exceeding 2^32-1.
+func.func @large_step_size(%A: memref<4294967295xf32>) {
+ %cst = arith.constant 0.000000e+00 : f32
+ // CHECK: affine.for %{{.*}} = 0 to 256 step 549755813760 {
+ affine.for %i = 0 to 256 step 4294967295 {
+   affine.store %cst, %A[%i] : memref<4294967295xf32>
+ }
+ return
+}
</pre>

</details>

https://github.com/llvm/llvm-project/pull/66026