Skip to content

[mlir][affine] Support vectorization with the step size exceeding 2^32-1 #66026

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][affine] Support vectorization with the step size exceeding 2^32-1 #66026

wants to merge 1 commit into from

Conversation

Lewuathe
Copy link
Member

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.

@Lewuathe Lewuathe requested a review from a team as a code owner September 11, 2023 23:07
@llvmbot
Copy link
Member

llvmbot commented Sep 11, 2023

@llvm/pr-subscribers-mlir-affine

Changes

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

@llvmbot
Copy link
Member

llvmbot commented Sep 11, 2023

@llvm/pr-subscribers-mlir

Changes

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

@llvmbot
Copy link
Member

llvmbot commented Sep 11, 2023

@llvm/pr-subscribers-mlir-core

Changes

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

sergei-grechanik
sergei-grechanik reviewed Sep 22, 2023
View reviewed changes
mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
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 {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The problem is that currently loops with non-unit step sizes are not vectorized correctly (we should probably add a check and avoid vectorizing these loops).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@sergei-grechanik sergei-grechanik sergei-grechanik left review comments

Assignees
No one assigned
Labels
mlir:affine mlir:core MLIR Core Infrastructure mlir
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@Lewuathe @llvmbot @sergei-grechanik @dcaballe