[Mlir-commits] [mlir] [mlir][vector] Allow signless integer element types in `vector.step` (PR #205142)
Victor Perez
llvmlistbot at llvm.org
Fri Jun 26 09:54:43 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/205142
>From 180ea5118ce6dd805251a4a688403aafb4a56c57 Mon Sep 17 00:00:00 2001
From: Victor Perez Carrasco <victorperez at meta.com>
Date: Mon, 22 Jun 2026 08:09:56 -0700
Subject: [PATCH] [mlir][vector] Allow signless integer element types in
vector.step
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
`vector.step` previously produced only `index` vectors. Allow its result
element type to also be a signless integer of at least 8 bits. If a lane's
value in the sequence is not representable in the element type it wraps
around (the value is truncated to the element bitwidth), consistent with
the `llvm.stepvector` intrinsic that `vector.step` lowers to.
Integer range inference and the VectorToSPIRV and VectorToLLVM lowerings
are updated to honor the new element types and the wrap-around semantics.
`populateVectorStepLoweringPatterns` takes an `indexBitwidth` controlling
how `index`-typed steps are materialized as constants, with 0 meaning
"leave them untouched"; VectorToLLVM uses 0 so that `index` steps keep
lowering through `llvm.stepvector`, whose result type carries the target
index bitwidth, instead of a target-independent constant.
The wrap-around path has no SPIR-V test: SPIR-V vector types are limited
to at most a handful of elements, so no legal SPIR-V vector is large
enough to exceed an 8-bit (or wider) element range.
Signed-off-by: Víctor Pérez Carrasco <victorperez at fb.com>
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 21 +++++++---
.../Vector/Transforms/LoweringPatterns.h | 6 ++-
.../VectorToLLVM/ConvertVectorToLLVM.cpp | 10 +----
.../VectorToLLVM/ConvertVectorToLLVMPass.cpp | 3 +-
.../VectorToSPIRV/VectorToSPIRV.cpp | 6 ++-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 12 ++++--
.../Vector/Transforms/LowerVectorStep.cpp | 39 ++++++++++++++---
.../vector-to-llvm-interface.mlir | 19 +++++++++
.../VectorToLLVM/vector-to-llvm.mlir | 28 +++++++++++--
.../VectorToSPIRV/vector-to-spirv.mlir | 39 +++++++++++++++++
.../Vector/canonicalize/vector-step.mlir | 39 +++++++++++++++++
.../Dialect/Vector/int-range-interface.mlir | 35 ++++++++++++++++
mlir/test/Dialect/Vector/invalid.mlir | 36 +++++++++++++++-
mlir/test/Dialect/Vector/ops.mlir | 6 +++
.../Dialect/Vector/vector-step-lowering.mlir | 42 +++++++++++++++++++
.../Dialect/Vector/TestVectorTransforms.cpp | 27 ++++++++++++
16 files changed, 334 insertions(+), 34 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/vector-step-lowering.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 7578ce78a0f00..a5ef4b4897894 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -3019,6 +3019,10 @@ def Vector_ScanOp :
// VectorStepOp
//===----------------------------------------------------------------------===//
+def VectorStepElementType : Type<
+ CPred<"::llvm::isa<::mlir::IndexType>($_self) || ($_self.isSignlessInteger() && $_self.getIntOrFloatBitWidth() >= 8)">,
+ "index or signless integer of at least 8 bits">;
+
def Vector_StepOp : Vector_Op<"step", [
Pure,
DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
@@ -3026,20 +3030,25 @@ def Vector_StepOp : Vector_Op<"step", [
]> {
let summary = "A linear sequence of values from 0 to N";
let description = [{
- A `step` operation produces an index vector, i.e. a 1-D vector of values of
- index type that represents a linear sequence from 0 to N-1, where N is the
- number of elements in the `result` vector.
+ A `step` operation produces a 1-D vector representing a linear sequence from
+ 0 to N-1, where N is the number of elements in the `result` vector.
+
+ The result element type must be `index` or a signless integer of at least 8
+ bits. If the sequence value exceeds the allowed limit for the element type
+ then the result for that lane is truncated.
Supports fixed-width and scalable vectors.
Examples:
```mlir
- %0 = vector.step : vector<4xindex> ; [0, 1, 2, 3]
- %1 = vector.step : vector<[4]xindex> ; [0, 1, .., <vscale * 4 - 1>]
+ %0 = vector.step : vector<4xindex> // [0, 1, 2, 3]
+ %1 = vector.step : vector<4xi32> // [0, 1, 2, 3]
+ %2 = vector.step : vector<258xi8> // [0, 1, .., 255, 0, 1]
+ %3 = vector.step : vector<[4]xindex> // [0, 1, .., <vscale * 4 - 1>]
```
}];
- let results = (outs VectorOfRankAndType<[1], [Index]>:$result);
+ let results = (outs VectorOfRankAndType<[1], [VectorStepElementType]>:$result);
let assemblyFormat = "attr-dict `:` type($result)";
let hasCanonicalizer = 1;
}
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index d23a4d5c3f5fb..4934713131503 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -241,8 +241,12 @@ void populateVectorScanLoweringPatterns(RewritePatternSet &patterns,
/// Populate the pattern set with the following patterns:
///
/// [StepToArithConstantOp]
-/// Convert vector.step op into arith ops if not using scalable vectors
+/// Convert a non-scalable `vector.step` into an `arith.constant`. `index`-typed
+/// steps are materialized using `indexBitwidth` as the index bitwidth;
+/// an `indexBitwidth` of 0 leaves them untouched. `indexBitwidth` must not
+/// exceed `IndexType::kInternalStorageBitWidth`.
void populateVectorStepLoweringPatterns(RewritePatternSet &patterns,
+ unsigned indexBitwidth = 64,
PatternBenefit benefit = 1);
/// Populate the pattern set with the following patterns:
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index 8e9d37648841a..eafb117894d28 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -2030,17 +2030,12 @@ struct VectorToElementsLowering
};
/// Conversion pattern for vector.step.
-struct VectorScalableStepOpLowering
- : public ConvertOpToLLVMPattern<vector::StepOp> {
+struct VectorStepOpLowering : public ConvertOpToLLVMPattern<vector::StepOp> {
using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
LogicalResult
matchAndRewrite(vector::StepOp stepOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- auto resultType = cast<VectorType>(stepOp.getType());
- if (!resultType.isScalable()) {
- return failure();
- }
Type llvmType = typeConverter->convertType(stepOp.getType());
rewriter.replaceOpWithNewOp<LLVM::StepVectorOp>(stepOp, llvmType);
return success();
@@ -2275,8 +2270,7 @@ void mlir::populateVectorToLLVMConversionPatterns(
VectorScalableInsertOpLowering, VectorScalableExtractOpLowering,
MaskedReductionOpConversion, VectorInterleaveOpLowering,
VectorDeinterleaveOpLowering, VectorFromElementsLowering,
- VectorToElementsLowering, VectorScalableStepOpLowering>(
- converter);
+ VectorToElementsLowering, VectorStepOpLowering>(converter);
}
namespace {
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
index 2b358d312dcfe..e3174829d90c0 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
@@ -87,7 +87,8 @@ void ConvertVectorToLLVMPass::runOnOperation() {
populateVectorMaskMaterializationPatterns(patterns,
force32BitVectorIndices);
populateVectorInsertExtractStridedSliceTransforms(patterns);
- populateVectorStepLoweringPatterns(patterns);
+ // Do not lower index vector.step to arith.constant.
+ populateVectorStepLoweringPatterns(patterns, /*indexBitwidth=*/0);
populateVectorRankReducingFMAPattern(patterns);
populateVectorGatherLoweringPatterns(patterns);
populateVectorFromElementsUnrollPatterns(patterns);
diff --git a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 423f2840851d4..c6be59471b6d5 100644
--- a/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
+++ b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
@@ -1024,8 +1024,10 @@ struct VectorStepOpConvert final : OpConversionPattern<vector::StepOp> {
Location loc = stepOp.getLoc();
int64_t numElements = stepOp.getType().getNumElements();
- auto intType =
- rewriter.getIntegerType(typeConverter.getIndexTypeBitwidth());
+ // Handle vector<1 x type> case, converting to the scalar type.
+ Type intType = isa<VectorType>(dstType)
+ ? cast<VectorType>(dstType).getElementType()
+ : dstType;
// Input vectors of size 1 are converted to scalars by the type converter.
// We just create a constant in this case.
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 67c31730f4b65..9d5043c6d8b48 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -8213,10 +8213,14 @@ void StepOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
return;
}
unsigned bitwidth = ConstantIntRanges::getStorageBitwidth(resultType);
- APInt zero(bitwidth, 0);
- APInt high(bitwidth, resultType.getDimSize(0) - 1);
- ConstantIntRanges result = {zero, high, zero, high};
- setResultRanges(getResult(), result);
+ // The result holds the sequence [0, 1, ..., N-1], with each value truncated
+ // to the result element type.
+ uint64_t maxIndex = resultType.getDimSize(0) - 1;
+ APInt umin = APInt::getZero(bitwidth);
+ APInt umax = APInt::getMaxValue(bitwidth).ugt(maxIndex)
+ ? APInt(bitwidth, maxIndex)
+ : APInt::getMaxValue(bitwidth);
+ setResultRanges(getResult(), ConstantIntRanges::fromUnsigned(umin, umax));
}
namespace {
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorStep.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorStep.cpp
index 08e7c895831ce..e87c6dbbf28a7 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorStep.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorStep.cpp
@@ -24,7 +24,12 @@ using namespace mlir::vector;
namespace {
struct StepToArithConstantOpRewrite final : OpRewritePattern<vector::StepOp> {
- using Base::Base;
+ StepToArithConstantOpRewrite(MLIRContext *context, unsigned indexBitwidth,
+ PatternBenefit benefit)
+ : OpRewritePattern(context, benefit), indexBitwidth(indexBitwidth) {
+ assert(indexBitwidth <= IndexType::kInternalStorageBitWidth &&
+ "indexBitwidth cannot exceed the index storage bitwidth");
+ }
LogicalResult matchAndRewrite(vector::StepOp stepOp,
PatternRewriter &rewriter) const override {
@@ -32,18 +37,40 @@ struct StepToArithConstantOpRewrite final : OpRewritePattern<vector::StepOp> {
if (resultType.isScalable()) {
return failure();
}
+ Type elementType = resultType.getElementType();
+ // An `indexBitwidth` of 0 means "leave `index`-typed steps alone": the
+ // index bitwidth is target-dependent, so callers that don't know it (and
+ // defer to a later lowering, e.g. to `llvm.intr.stepvector`) opt out.
+ if (elementType.isIndex() && indexBitwidth == 0) {
+ return failure();
+ }
+ // Values wrap around at `computeWidth`. `index` elements are stored in a
+ // `DenseElementsAttr` using the internal storage bitwidth, so the wrapped
+ // value is widened to it; integer elements use their own bitwidth.
+ unsigned computeWidth = elementType.isIndex()
+ ? indexBitwidth
+ : elementType.getIntOrFloatBitWidth();
+ unsigned storageWidth = elementType.isIndex()
+ ? IndexType::kInternalStorageBitWidth
+ : computeWidth;
int64_t elementCount = resultType.getNumElements();
- SmallVector<APInt> indices =
- llvm::map_to_vector(llvm::seq(elementCount),
- [](int64_t i) { return APInt(/*width=*/64, i); });
+ SmallVector<APInt> indices = llvm::map_to_vector(
+ llvm::seq(elementCount), [computeWidth, storageWidth](int64_t i) {
+ return APInt(computeWidth, i, /*isSigned=*/false,
+ /*implicitTrunc=*/true)
+ .zext(storageWidth);
+ });
rewriter.replaceOpWithNewOp<arith::ConstantOp>(
stepOp, DenseElementsAttr::get(resultType, indices));
return success();
}
+
+ unsigned indexBitwidth;
};
} // namespace
void mlir::vector::populateVectorStepLoweringPatterns(
- RewritePatternSet &patterns, PatternBenefit benefit) {
- patterns.add<StepToArithConstantOpRewrite>(patterns.getContext(), benefit);
+ RewritePatternSet &patterns, unsigned indexBitwidth, PatternBenefit benefit) {
+ patterns.add<StepToArithConstantOpRewrite>(patterns.getContext(),
+ indexBitwidth, benefit);
}
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
index e82f75d93066a..2839a039959fa 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
@@ -2189,6 +2189,17 @@ func.func @to_elements_dead_elements(%a: vector<4xf32>) -> (f32, f32) {
// vector.step
//===----------------------------------------------------------------------===//
+// Fixed-width `index` steps lower to `llvm.intr.stepvector` on both the
+// `--convert-to-llvm` and `-convert-vector-to-llvm` paths.
+// CHECK-LABEL: @step
+// CHECK: %[[STEPVECTOR:.*]] = llvm.intr.stepvector : vector<4xi64>
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[STEPVECTOR]] : vector<4xi64> to vector<4xindex>
+// CHECK: return %[[CAST]] : vector<4xindex>
+func.func @step() -> vector<4xindex> {
+ %0 = vector.step : vector<4xindex>
+ return %0 : vector<4xindex>
+}
+
// CHECK-LABEL: @step_scalable
// CHECK: %[[STEPVECTOR:.*]] = llvm.intr.stepvector : vector<[4]xi64>
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[STEPVECTOR]] : vector<[4]xi64> to vector<[4]xindex>
@@ -2197,3 +2208,11 @@ func.func @step_scalable() -> vector<[4]xindex> {
%0 = vector.step : vector<[4]xindex>
return %0 : vector<[4]xindex>
}
+
+// CHECK-LABEL: @step_scalable_i8
+// CHECK: %[[STEPVECTOR:.*]] = llvm.intr.stepvector : vector<[4]xi8>
+// CHECK: return %[[STEPVECTOR]] : vector<[4]xi8>
+func.func @step_scalable_i8() -> vector<[4]xi8> {
+ %0 = vector.step : vector<[4]xi8>
+ return %0 : vector<[4]xi8>
+}
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
index 77f60b3172296..8335e68b57430 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
@@ -1746,16 +1746,36 @@ func.func @deinterleave_2d_scalable(%arg: vector<2x[8]xf32>) -> (vector<2x[4]xf3
// vector.step
//===----------------------------------------------------------------------===//
-// TODO: Investigate why this wouldn't lower with --convert-to-llvm="filter-dialects=vector"
-
+// `index`-typed `vector.step` is left for the `llvm.intr.stepvector` lowering
+// (the index bitwidth is target-dependent); fixed-width integers below lower to
+// an `arith.constant`.
// CHECK-LABEL: @step
-// CHECK: %[[CST:.+]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
-// CHECK: return %[[CST]] : vector<4xindex>
+// CHECK: %[[STEPVECTOR:.*]] = llvm.intr.stepvector : vector<4xi64>
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[STEPVECTOR]] : vector<4xi64> to vector<4xindex>
+// CHECK: return %[[CAST]] : vector<4xindex>
func.func @step() -> vector<4xindex> {
%0 = vector.step : vector<4xindex>
return %0 : vector<4xindex>
}
+// CHECK-LABEL: @step_i8
+// CHECK: %[[CST:.+]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xi8>
+// CHECK: return %[[CST]] : vector<4xi8>
+func.func @step_i8() -> vector<4xi8> {
+ %0 = vector.step : vector<4xi8>
+ return %0 : vector<4xi8>
+}
+
+// The sequence wraps past the i8 limit, so lane values are truncated. The
+// constant is printed as a hex blob ending in `...FEFF0001`: after 255 (`FF`)
+// the values restart at 0 (`00`) and 1 (`01`).
+// CHECK-LABEL: @step_i8_truncate
+// CHECK: %[[CST:.+]] = arith.constant dense<"0x0001{{.*}}FEFF0001"> : vector<258xi8>
+// CHECK: return %[[CST]] : vector<258xi8>
+func.func @step_i8_truncate() -> vector<258xi8> {
+ %0 = vector.step : vector<258xi8>
+ return %0 : vector<258xi8>
+}
// -----
diff --git a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
index 152d32a81133d..f904dd9d35c37 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -1065,6 +1065,45 @@ func.func @step_size1() -> vector<1xindex> {
// -----
+module attributes {
+ spirv.target_env = #spirv.target_env<
+ #spirv.vce<v1.0, [Int8, Shader], []>, #spirv.resource_limits<>>
+ } {
+
+// CHECK-LABEL: @step_i8()
+// CHECK: %[[CST0:.*]] = spirv.Constant 0 : i8
+// CHECK: %[[CST1:.*]] = spirv.Constant 1 : i8
+// CHECK: %[[CST2:.*]] = spirv.Constant 2 : i8
+// CHECK: %[[CST3:.*]] = spirv.Constant 3 : i8
+// CHECK: %[[CONSTRUCT:.*]] = spirv.CompositeConstruct %[[CST0]], %[[CST1]], %[[CST2]], %[[CST3]] : (i8, i8, i8, i8) -> vector<4xi8>
+// CHECK: return %[[CONSTRUCT]] : vector<4xi8>
+func.func @step_i8() -> vector<4xi8> {
+ %0 = vector.step : vector<4xi8>
+ return %0 : vector<4xi8>
+}
+
+}
+
+// -----
+
+module attributes {
+ spirv.target_env = #spirv.target_env<
+ #spirv.vce<v1.0, [Int16, Shader], []>, #spirv.resource_limits<>>
+ } {
+
+// CHECK-LABEL: @step_size1_i16()
+// CHECK: %[[CST0:.*]] = spirv.Constant 0 : i16
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[CST0]] : i16 to vector<1xi16>
+// CHECK: return %[[CAST]] : vector<1xi16>
+func.func @step_size1_i16() -> vector<1xi16> {
+ %0 = vector.step : vector<1xi16>
+ return %0 : vector<1xi16>
+}
+
+}
+
+// -----
+
module attributes {
spirv.target_env = #spirv.target_env<
#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-step.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-step.mlir
index 023a0e52b65dc..7592e994449b7 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-step.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-step.mlir
@@ -309,3 +309,42 @@ func.func @negative_ne_constant_2() -> vector<3xi1> {
return %1 : vector<3xi1>
}
+// -----
+
+// CHECK-LABEL: @ult_i32_constant_4_rhs
+// CHECK: %[[CST:.*]] = arith.constant dense<true> : vector<4xi1>
+// CHECK: return %[[CST]] : vector<4xi1>
+func.func @ult_i32_constant_4_rhs() -> vector<4xi1> {
+ %cst = arith.constant dense<4> : vector<4xi32>
+ %0 = vector.step : vector<4xi32>
+ // [0, 1, 2, 3] < 4 => [true, true, true, true] => fold
+ %1 = arith.cmpi ult, %0, %cst : vector<4xi32>
+ return %1 : vector<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @negative_ult_i32_constant_2_rhs
+// CHECK: %[[CMP:.*]] = arith.cmpi
+// CHECK: return %[[CMP]]
+func.func @negative_ult_i32_constant_2_rhs() -> vector<4xi1> {
+ %cst = arith.constant dense<2> : vector<4xi32>
+ %0 = vector.step : vector<4xi32>
+ // [0, 1, 2, 3] < 2 => [true, true, false, false] => don't fold
+ %1 = arith.cmpi ult, %0, %cst : vector<4xi32>
+ return %1 : vector<4xi1>
+}
+
+// -----
+
+// The element count (300) is greater than the constant (100), so the
+// count-based fold guard (`stepSize <= constValue`) already blocks the fold.
+// CHECK-LABEL: @negative_ult_i8_wrap
+// CHECK: %[[CMP:.*]] = arith.cmpi
+// CHECK: return %[[CMP]]
+func.func @negative_ult_i8_wrap() -> vector<300xi1> {
+ %cst = arith.constant dense<100> : vector<300xi8>
+ %0 = vector.step : vector<300xi8>
+ %1 = arith.cmpi ult, %0, %cst : vector<300xi8>
+ return %1 : vector<300xi1>
+}
diff --git a/mlir/test/Dialect/Vector/int-range-interface.mlir b/mlir/test/Dialect/Vector/int-range-interface.mlir
index 4da8d8a967c73..6e2aa064a67f8 100644
--- a/mlir/test/Dialect/Vector/int-range-interface.mlir
+++ b/mlir/test/Dialect/Vector/int-range-interface.mlir
@@ -116,3 +116,38 @@ func.func @vector_step() -> vector<8xindex> {
%1 = test.reflect_bounds %0 : vector<8xindex>
func.return %1 : vector<8xindex>
}
+
+// CHECK-LABEL: func @vector_step_i32
+// CHECK: test.reflect_bounds {smax = 7 : si32, smin = 0 : si32, umax = 7 : ui32, umin = 0 : ui32}
+func.func @vector_step_i32() -> vector<8xi32> {
+ %0 = vector.step : vector<8xi32>
+ %1 = test.reflect_bounds %0 : vector<8xi32>
+ func.return %1 : vector<8xi32>
+}
+
+// Boundary: 255 lanes do not wrap, so the upper bound is the last lane (254).
+// CHECK-LABEL: func @vector_step_i8_no_wrap_boundary
+// CHECK: test.reflect_bounds {smax = 127 : si8, smin = -128 : si8, umax = 254 : ui8, umin = 0 : ui8}
+func.func @vector_step_i8_no_wrap_boundary() -> vector<255xi8> {
+ %0 = vector.step : vector<255xi8>
+ %1 = test.reflect_bounds %0 : vector<255xi8>
+ func.return %1 : vector<255xi8>
+}
+
+// Boundary: 256 lanes exactly fill the i8 range without wrapping (last lane 255).
+// CHECK-LABEL: func @vector_step_i8_exact_fit
+// CHECK: test.reflect_bounds {smax = 127 : si8, smin = -128 : si8, umax = 255 : ui8, umin = 0 : ui8}
+func.func @vector_step_i8_exact_fit() -> vector<256xi8> {
+ %0 = vector.step : vector<256xi8>
+ %1 = test.reflect_bounds %0 : vector<256xi8>
+ func.return %1 : vector<256xi8>
+}
+
+// The sequence wraps (300 > 256), so the result spans the entire i8 range.
+// CHECK-LABEL: func @vector_step_i8_wrap
+// CHECK: test.reflect_bounds {smax = 127 : si8, smin = -128 : si8, umax = 255 : ui8, umin = 0 : ui8}
+func.func @vector_step_i8_wrap() -> vector<300xi8> {
+ %0 = vector.step : vector<300xi8>
+ %1 = test.reflect_bounds %0 : vector<300xi8>
+ func.return %1 : vector<300xi8>
+}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 2fed3002596a3..e7931c49cd1c3 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -2056,7 +2056,7 @@ func.func @invalid_from_elements_scalable(%a: f32, %b: i32) {
// -----
func.func @invalid_step_0d() {
- // expected-error @+1 {{vector.step' op result #0 must be vector of index values of ranks 1, but got 'vector<f32>'}}
+ // expected-error @+1 {{vector.step' op result #0 must be vector of index or signless integer of at least 8 bits values of ranks 1, but got 'vector<f32>'}}
vector.step : vector<f32>
return
}
@@ -2064,13 +2064,45 @@ func.func @invalid_step_0d() {
// -----
func.func @invalid_step_2d() {
- // expected-error @+1 {{vector.step' op result #0 must be vector of index values of ranks 1, but got 'vector<2x4xf32>'}}
+ // expected-error @+1 {{vector.step' op result #0 must be vector of index or signless integer of at least 8 bits values of ranks 1, but got 'vector<2x4xf32>'}}
vector.step : vector<2x4xf32>
return
}
// -----
+func.func @invalid_step_float_element() {
+ // expected-error @+1 {{vector.step' op result #0 must be vector of index or signless integer of at least 8 bits values of ranks 1, but got 'vector<4xf32>'}}
+ vector.step : vector<4xf32>
+ return
+}
+
+// -----
+
+func.func @invalid_step_narrow_integer() {
+ // expected-error @+1 {{vector.step' op result #0 must be vector of index or signless integer of at least 8 bits values of ranks 1, but got 'vector<4xi4>'}}
+ vector.step : vector<4xi4>
+ return
+}
+
+// -----
+
+func.func @invalid_step_i1_element() {
+ // expected-error @+1 {{vector.step' op result #0 must be vector of index or signless integer of at least 8 bits values of ranks 1, but got 'vector<4xi1>'}}
+ vector.step : vector<4xi1>
+ return
+}
+
+// -----
+
+func.func @invalid_step_unsigned_integer() {
+ // expected-error @+1 {{vector.step' op result #0 must be vector of index or signless integer of at least 8 bits values of ranks 1, but got 'vector<4xui8>'}}
+ vector.step : vector<4xui8>
+ return
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// vector.load
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index de620221944de..e84bd3f1dce17 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -1158,6 +1158,12 @@ func.func @step() {
%0 = vector.step : vector<2xindex>
// CHECK: vector.step : vector<[4]xindex>
%1 = vector.step : vector<[4]xindex>
+ // CHECK: vector.step : vector<2xi32>
+ %2 = vector.step : vector<2xi32>
+ // CHECK: vector.step : vector<8xi8>
+ %3 = vector.step : vector<8xi8>
+ // CHECK: vector.step : vector<[4]xi16>
+ %4 = vector.step : vector<[4]xi16>
return
}
diff --git a/mlir/test/Dialect/Vector/vector-step-lowering.mlir b/mlir/test/Dialect/Vector/vector-step-lowering.mlir
new file mode 100644
index 0000000000000..ac7b894a52175
--- /dev/null
+++ b/mlir/test/Dialect/Vector/vector-step-lowering.mlir
@@ -0,0 +1,42 @@
+// RUN: mlir-opt %s -test-vector-step-lowering=index-bitwidth=64 -split-input-file | FileCheck %s --check-prefixes=CHECK,WIDTH64
+// RUN: mlir-opt %s -test-vector-step-lowering -split-input-file | FileCheck %s --check-prefixes=CHECK,SKIP
+// RUN: mlir-opt %s -test-vector-step-lowering=index-bitwidth=8 -split-input-file | FileCheck %s --check-prefixes=CHECK,WIDTH8
+
+// `index`-typed steps are materialized with `index-bitwidth` as the index
+// bitwidth; a bitwidth of 0 (the default) leaves them untouched. Integer steps
+// are unaffected by the option.
+
+// CHECK-LABEL: @step_index
+// WIDTH64: arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
+// SKIP: vector.step : vector<4xindex>
+// WIDTH8: arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
+func.func @step_index() -> vector<4xindex> {
+ %0 = vector.step : vector<4xindex>
+ return %0 : vector<4xindex>
+}
+
+// -----
+
+// The lane values only wrap when the index bitwidth is small enough. With
+// `index-bitwidth=8` the sequence wraps at 256, so element 256 becomes 0 and
+// element 257 becomes 1; with the default 64-bit index it does not wrap (256
+// and 257 are represented as-is).
+
+// CHECK-LABEL: @step_index_wrap
+// WIDTH64: arith.constant dense<"0x0000000000000000{{.*}}00010000000000000101000000000000"> : vector<258xindex>
+// SKIP: vector.step : vector<258xindex>
+// WIDTH8: arith.constant dense<"0x0000000000000000{{.*}}00000000000000000100000000000000"> : vector<258xindex>
+func.func @step_index_wrap() -> vector<258xindex> {
+ %0 = vector.step : vector<258xindex>
+ return %0 : vector<258xindex>
+}
+
+// -----
+
+// Integer steps are independent of `index-bitwidth`.
+// CHECK-LABEL: @step_i8
+// CHECK: arith.constant dense<[0, 1, 2, 3]> : vector<4xi8>
+func.func @step_i8() -> vector<4xi8> {
+ %0 = vector.step : vector<4xi8>
+ return %0 : vector<4xi8>
+}
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index 043181c16c759..971b2f3c1e7af 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -545,6 +545,31 @@ struct TestVectorScanLowering
}
};
+struct TestVectorStepLowering
+ : public PassWrapper<TestVectorStepLowering, OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorStepLowering)
+
+ TestVectorStepLowering() = default;
+ TestVectorStepLowering(const TestVectorStepLowering &pass)
+ : PassWrapper(pass) {}
+
+ StringRef getArgument() const final { return "test-vector-step-lowering"; }
+ StringRef getDescription() const final {
+ return "Test lowering patterns that lower the step op in the vector "
+ "dialect";
+ }
+ Option<unsigned> indexBitwidth{
+ *this, "index-bitwidth",
+ llvm::cl::desc("Bitwidth used to materialize `index`-typed `vector.step` "
+ "(0 leaves them untouched)"),
+ llvm::cl::init(0)};
+ void runOnOperation() override {
+ RewritePatternSet patterns(&getContext());
+ populateVectorStepLoweringPatterns(patterns, indexBitwidth);
+ (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+ }
+};
+
/// Allocate shared memory for a single warp to test lowering of
/// WarpExecuteOnLane0Op.
static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
@@ -1091,6 +1116,8 @@ void registerTestVectorLowerings() {
PassRegistration<TestVectorScanLowering>();
+ PassRegistration<TestVectorStepLowering>();
+
PassRegistration<TestVectorShuffleLowering>();
PassRegistration<TestVectorDistribution>();
More information about the Mlir-commits
mailing list