[Mlir-commits] [mlir] Allow signless integer element types in vector.step (PR #205142)
Victor Perez
llvmlistbot at llvm.org
Mon Jun 22 09:58:04 PDT 2026
https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/205142
`vector.step` previously produced only `index` vectors. Extend it to also accept a signless integer element type of at least 8 bits. Also clarify wraparound behaviour: if a sequence value exceeds the limit of the element type, the result for that lane is truncated. This is backwards compatible with the lowering of `vector.step` to the `llvm.stepvector` intrinsic.
- Op definition: the result element type is now `index` or a signless integer of at least 8 bits; the description documents the truncation behavior.
- LowerVectorStep: build the constant using the element type's bitwidth so out-of-range lanes wrap around.
- StepOp::inferResultRanges: clamp the unsigned upper bound and derive the signed bounds via ConstantIntRanges::fromUnsigned, so wrapping sequences are modeled correctly.
- VectorToSPIRV: materialize the constants in the result element type rather than the index bitwidth.
- Add tests for parsing/printing, verification, canonicalization, integer range inference, and the VectorToLLVM / VectorToSPIRV lowerings (including a truncating case).
>From e7d212cb71004feffc5d3e0ebd7fe84c4a469109 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] Allow signless integer element types in vector.step
vector.step previously produced only `index` vectors. Extend it to also
accept a signless integer element type of at least 8 bits, matching the
`llvm.stepvector` intrinsic that the scalable case lowers to. Following the
same semantics as `llvm.stepvector`: if a sequence value exceeds the limit of
the element type, the result for that lane is truncated.
- Op definition: the result element type is now `index` or a signless integer
of at least 8 bits; the description documents the truncation behavior.
- LowerVectorStep: build the constant using the element type's bitwidth so
out-of-range lanes wrap around.
- StepOp::inferResultRanges: clamp the unsigned upper bound and derive the
signed bounds via ConstantIntRanges::fromUnsigned, so wrapping sequences are
modeled correctly.
- VectorToSPIRV: materialize the constants in the result element type rather
than the index bitwidth.
- Add tests for parsing/printing, verification, canonicalization, integer
range inference, and the VectorToLLVM / VectorToSPIRV lowerings (including a
truncating case).
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 25 +++++++----
.../VectorToSPIRV/VectorToSPIRV.cpp | 6 ++-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 14 +++++--
.../Vector/Transforms/LowerVectorStep.cpp | 15 +++++--
.../vector-to-llvm-interface.mlir | 8 ++++
.../VectorToLLVM/vector-to-llvm.mlir | 19 +++++++++
.../VectorToSPIRV/vector-to-spirv.mlir | 42 +++++++++++++++++++
.../Vector/canonicalize/vector-step.mlir | 25 +++++++++++
.../Dialect/Vector/int-range-interface.mlir | 17 ++++++++
mlir/test/Dialect/Vector/invalid.mlir | 36 +++++++++++++++-
mlir/test/Dialect/Vector/ops.mlir | 6 +++
11 files changed, 195 insertions(+), 18 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 7578ce78a0f00..f495e9d6f7f55 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1465,7 +1465,7 @@ def Vector_TransferReadOp :
"AffineMapAttr":$permutationMapAttr,
"ArrayAttr":$inBoundsAttr)>,
/// 2. Builder that sets padding to `padding` or poison if not provided and
- /// an empty mask (variant without attrs).
+ /// an empty mask (variant without attrs).
/// If `padding` is null, a poison value is used.
/// If `permutationMap` is null, a minor identity map is used.
/// If `inBounds` is null, an empty mask is used.
@@ -3019,6 +3019,12 @@ def Vector_ScanOp :
// VectorStepOp
//===----------------------------------------------------------------------===//
+// Allowed element type for `vector.step`: `index`, or a signless integer of at
+// least 8 bits.
+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 +3032,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/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp b/mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
index 423f2840851d4..c2b4b37d7f9ae 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 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..4ea7e0ec7a64f 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -8213,10 +8213,16 @@ 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. If `N - 1` is not representable the sequence
+ // wraps and spans every value, so clamp the unsigned upper bound to the
+ // largest representable value.
+ 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..808b2623404a9 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorStep.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorStep.cpp
@@ -33,9 +33,18 @@ struct StepToArithConstantOpRewrite final : OpRewritePattern<vector::StepOp> {
return failure();
}
int64_t elementCount = resultType.getNumElements();
- SmallVector<APInt> indices =
- llvm::map_to_vector(llvm::seq(elementCount),
- [](int64_t i) { return APInt(/*width=*/64, i); });
+ Type elementType = resultType.getElementType();
+ // `index` elements are stored in a `DenseElementsAttr` as 64-bit values;
+ // integer elements use their own bitwidth. Values that are not
+ // representable in the element type wrap around (i.e. lane `i` holds `i`
+ // truncated to the element type).
+ unsigned bitWidth = elementType.isIndex()
+ ? IndexType::kInternalStorageBitWidth
+ : elementType.getIntOrFloatBitWidth();
+ SmallVector<APInt> indices = llvm::map_to_vector(
+ llvm::seq(elementCount), [bitWidth](int64_t i) {
+ return APInt(bitWidth, i, /*isSigned=*/false, /*implicitTrunc=*/true);
+ });
rewriter.replaceOpWithNewOp<arith::ConstantOp>(
stepOp, DenseElementsAttr::get(resultType, indices));
return success();
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
index e82f75d93066a..98bb2eafeafab 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
@@ -2197,3 +2197,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..9cb9b1f5ab730 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
@@ -1756,6 +1756,25 @@ func.func @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..837185c430efc 100644
--- a/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
+++ b/mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
@@ -1065,6 +1065,48 @@ 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, [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, [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..0b704f4a39c1a 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-step.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-step.mlir
@@ -309,3 +309,28 @@ 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>
+}
diff --git a/mlir/test/Dialect/Vector/int-range-interface.mlir b/mlir/test/Dialect/Vector/int-range-interface.mlir
index 4da8d8a967c73..ae42cc3a4657c 100644
--- a/mlir/test/Dialect/Vector/int-range-interface.mlir
+++ b/mlir/test/Dialect/Vector/int-range-interface.mlir
@@ -116,3 +116,20 @@ 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>
+}
+
+// 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
}
More information about the Mlir-commits
mailing list