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

[Kai Sasaki]

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

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. 

>From 165104cbeb8ba50a50c0aca270756efb63afc43d Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Mon, 11 Sep 2023 16:32:47 +0900
Subject: [PATCH] [mlir][affine] Support vectorization with the step size
 exceeding 2^32-1

---
 .../Dialect/Affine/Transforms/SuperVectorize.cpp    |  2 +-
 .../Dialect/Affine/SuperVectorize/vectorize_1d.mlir | 13 +++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

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
+}