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

[Lewuathe]

Summary

AffineFor step size is capable of storing 2^64-1 as it retains int64_t. But due to the conversion to lower bit-width type (unsinged), it fails to calculate the right step size after vectorization. New step size also should keep the same size as the original step size.

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

LLVM_DEBUG(
         << "\n[early-vect]+++++ unsupported step size for reduction loop: "
         << forOp.getStep() << "\n");
     return nullptr;
   }
 
   // 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");
     int64_t forOpVecFactor = strategy.vectorSizes[vectorDim];
     newStep = forOp.getStep() * forOpVecFactor;
   } else {
     newStep = forOp.getStep();
   }

mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

affine.for %i = 0 to 256 {
    }
    affine.store %final_red, %out[%i] : memref<256xf32>
  }
  return
 }
 
 // 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
+}
 No newline at end of file