[Mlir-commits] [mlir] [mlir][vector] Add i8 to i2 truncation to the narrow type rewrites (PR #221185)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Sun Nov 1 05:27:54 PST 2026
https://github.com/Tim096 updated https://github.com/llvm/llvm-project/pull/221185
>From a00b482bb9bc76061444c4534577d225d679fbc3 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Fri, 4 Sep 2026 18:02:16 +0800
Subject: [PATCH 1/2] [mlir][vector] Add i8 to i2 truncation to the narrow type
rewrites
`RewriteAlignedSubByteIntTrunc` rewrites `arith.trunci` to `i4` and bails
on `i2`, while the extension side has handled both widths since #121298.
Add the `i8 -> i2` half: deinterleave twice to group the bytes by their
position in the destination byte, mask each to its low two bits, shift
into place, merge, and bitcast to `i2`. The `i8` truncation of a wider
source and the alignment preconditions are unchanged.
The `aligned_trunci_i8_to_i2_no_match` test becomes a positive test, and
the integration test gains an `i2` truncation printed as bits.
---
.../Transforms/VectorEmulateNarrowType.cpp | 68 ++++++++++++++++---
...vector-rewrite-subbyte-ext-and-trunci.mlir | 56 ++++++++++++++-
.../Vector/CPU/rewrite-narrow-types.mlir | 26 +++++++
3 files changed, 138 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index e237ae0568ae43..c343ca92b7e0b5 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -1976,6 +1976,57 @@ static Value rewriteI8ToI4Trunc(PatternRewriter &rewriter, Location loc,
return vector::BitCastOp::create(rewriter, loc, i4VecType, mergedHiLowOp);
}
+/// Rewrite the i8 -> i2 truncation into two levels of deinterleave and a
+/// series of bitwise ops to avoid leaving LLVM to scramble with peephole
+/// optimizations.
+static Value rewriteI8ToI2Trunc(PatternRewriter &rewriter, Location loc,
+ Value srcValue) {
+ VectorType srcVecType = cast<VectorType>(srcValue.getType());
+ assert(srcVecType.getElementType().isSignlessInteger(8) &&
+ "Expected i8 type");
+
+ // 1. De-interleave twice to group the i8 elements by their position within
+ // the destination byte.
+ // in = [0,1,2,3,4,5,6,7],...
+ // evenOdd = [0,2,4,6],... [1,3,5,7],...
+ // vec0 = [0,4],...
+ // vec1 = [1,5],...
+ // vec2 = [2,6],...
+ // vec3 = [3,7],...
+ auto evenOdd = vector::DeinterleaveOp::create(rewriter, loc, srcValue);
+ auto vec02 = vector::DeinterleaveOp::create(rewriter, loc, evenOdd.getRes1());
+ auto vec13 = vector::DeinterleaveOp::create(rewriter, loc, evenOdd.getRes2());
+
+ // 2. Zero out the upper side of each i8 element, then move the low i2 to
+ // its position in the byte. Position 3 needs no mask, the upper bits are
+ // shifted out.
+ VectorType deinterI8VecType = vec02.getResultVectorType();
+ constexpr int8_t i8LowBitMask = 0x03;
+ Value zeroOutMask = arith::ConstantOp::create(
+ rewriter, loc, DenseElementsAttr::get(deinterI8VecType, i8LowBitMask));
+ auto shiftLeft = [&](Value value, int8_t bitsToShift) {
+ Value shiftValues = arith::ConstantOp::create(
+ rewriter, loc, DenseElementsAttr::get(deinterI8VecType, bitsToShift));
+ return arith::ShLIOp::create(rewriter, loc, value, shiftValues);
+ };
+ Value elem0 =
+ arith::AndIOp::create(rewriter, loc, vec02.getRes1(), zeroOutMask);
+ Value elem1 = shiftLeft(
+ arith::AndIOp::create(rewriter, loc, vec13.getRes1(), zeroOutMask), 2);
+ Value elem2 = shiftLeft(
+ arith::AndIOp::create(rewriter, loc, vec02.getRes2(), zeroOutMask), 4);
+ Value elem3 = shiftLeft(vec13.getRes2(), 6);
+
+ // 3. Merge the four i2 values.
+ Value merged = arith::OrIOp::create(rewriter, loc, elem0, elem1);
+ merged = arith::OrIOp::create(rewriter, loc, merged, elem2);
+ merged = arith::OrIOp::create(rewriter, loc, merged, elem3);
+
+ // 4. Generate a bitcast vector<Xxi8> -> vector<4Xxi2>.
+ auto i2VecType = srcVecType.cloneWith(std::nullopt, rewriter.getI2Type());
+ return vector::BitCastOp::create(rewriter, loc, i2VecType, merged);
+}
+
namespace {
/// Rewrite bitcast(trunci) to a sequence of shuffles and bitwise ops that take
/// advantage of high-level information to avoid leaving LLVM to scramble with
@@ -2173,9 +2224,9 @@ struct RewriteAlignedSubByteIntExt : OpRewritePattern<ConversionOpType> {
}
};
-/// Rewrite the i8 -> i4 part of any truncation into a deinterleave and
-/// bitwise ops that take advantage of high-level information to avoid leaving
-/// LLVM to scramble with peephole optimizations.
+/// Rewrite the i8 -> i4 or i8 -> i2 part of any truncation into deinterleaves
+/// and bitwise ops that take advantage of high-level information to avoid
+/// leaving LLVM to scramble with peephole optimizations.
///
/// For example:
/// arith.trunci %in : vector<8xi32> to vector<8xi4>
@@ -2205,10 +2256,6 @@ struct RewriteAlignedSubByteIntTrunc : OpRewritePattern<arith::TruncIOp> {
if (failed(commonConversionPrecondition(rewriter, srcVecType, truncOp)))
return failure();
- // TODO: Add support for truncating to i2.
- if (dstVecType.getElementType().getIntOrFloatBitWidth() == 2)
- return failure();
-
// Check general alignment preconditions. We invert the src/dst type order
// to reuse the existing precondition logic.
if (failed(alignedConversionPrecondition(
@@ -2224,8 +2271,11 @@ struct RewriteAlignedSubByteIntTrunc : OpRewritePattern<arith::TruncIOp> {
? srcValue
: arith::TruncIOp::create(rewriter, loc, i8VecType, srcValue);
- // Rewrite the i8 -> i4 truncation part.
- Value subByteTrunc = rewriteI8ToI4Trunc(rewriter, loc, i8TruncVal);
+ // Rewrite the i8 -> i4 / i8 -> i2 truncation part. The alignment
+ // precondition above only admits these two widths.
+ Value subByteTrunc = dstVecType.getElementTypeBitWidth() == 2
+ ? rewriteI8ToI2Trunc(rewriter, loc, i8TruncVal)
+ : rewriteI8ToI4Trunc(rewriter, loc, i8TruncVal);
// Finalize the rewrite.
rewriter.replaceOp(truncOp, subByteTrunc);
diff --git a/mlir/test/Dialect/Vector/vector-rewrite-subbyte-ext-and-trunci.mlir b/mlir/test/Dialect/Vector/vector-rewrite-subbyte-ext-and-trunci.mlir
index aa75e020052566..1925bd812bddbb 100644
--- a/mlir/test/Dialect/Vector/vector-rewrite-subbyte-ext-and-trunci.mlir
+++ b/mlir/test/Dialect/Vector/vector-rewrite-subbyte-ext-and-trunci.mlir
@@ -203,13 +203,63 @@ func.func @aligned_trunci_nd(%a: vector<3x8x32xi32>) -> vector<3x8x32xi4> {
return %0 : vector<3x8x32xi4>
}
-func.func @aligned_trunci_i8_to_i2_no_match(%a: vector<8xi8>) -> vector<8xi2> {
- // CHECK-NOT: arith.bitcast
- // CHECK: arith.trunci %[[IN:.*]] : vector<8xi8> to vector<8xi2>
+// CHECK-LABEL: func.func @aligned_trunci_i8_to_i2(
+func.func @aligned_trunci_i8_to_i2(%a: vector<8xi8>) -> vector<8xi2> {
+// CHECK-SAME: %[[IN:.*]]: vector<8xi8>) -> vector<8xi2> {
+// CHECK-DAG: %[[LOW_MASK:.*]] = arith.constant dense<3> : vector<2xi8>
+// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<2> : vector<2xi8>
+// CHECK-DAG: %[[CST_4:.*]] = arith.constant dense<4> : vector<2xi8>
+// CHECK-DAG: %[[CST_6:.*]] = arith.constant dense<6> : vector<2xi8>
+// CHECK: %[[EVEN:.*]], %[[ODD:.*]] = vector.deinterleave %[[IN]] : vector<8xi8> -> vector<4xi8>
+// CHECK: %[[ELEM0:.*]], %[[ELEM2:.*]] = vector.deinterleave %[[EVEN]] : vector<4xi8> -> vector<2xi8>
+// CHECK: %[[ELEM1:.*]], %[[ELEM3:.*]] = vector.deinterleave %[[ODD]] : vector<4xi8> -> vector<2xi8>
+// CHECK: %[[ZEROED0:.*]] = arith.andi %[[ELEM0]], %[[LOW_MASK]] : vector<2xi8>
+// CHECK: %[[ZEROED1:.*]] = arith.andi %[[ELEM1]], %[[LOW_MASK]] : vector<2xi8>
+// CHECK: %[[SHL1:.*]] = arith.shli %[[ZEROED1]], %[[CST_2]] : vector<2xi8>
+// CHECK: %[[ZEROED2:.*]] = arith.andi %[[ELEM2]], %[[LOW_MASK]] : vector<2xi8>
+// CHECK: %[[SHL2:.*]] = arith.shli %[[ZEROED2]], %[[CST_4]] : vector<2xi8>
+// CHECK: %[[SHL3:.*]] = arith.shli %[[ELEM3]], %[[CST_6]] : vector<2xi8>
+// CHECK: %[[MERGED01:.*]] = arith.ori %[[ZEROED0]], %[[SHL1]] : vector<2xi8>
+// CHECK: %[[MERGED012:.*]] = arith.ori %[[MERGED01]], %[[SHL2]] : vector<2xi8>
+// CHECK: %[[MERGED:.*]] = arith.ori %[[MERGED012]], %[[SHL3]] : vector<2xi8>
+// CHECK: %[[I2:.*]] = vector.bitcast %[[MERGED]] : vector<2xi8> to vector<8xi2>
%0 = arith.trunci %a : vector<8xi8> to vector<8xi2>
return %0 : vector<8xi2>
}
+// CHECK-LABEL: func.func @aligned_trunci_i32_to_i2(
+func.func @aligned_trunci_i32_to_i2(%a: vector<8xi32>) -> vector<8xi2> {
+// CHECK-SAME: %[[IN:.*]]: vector<8xi32>) -> vector<8xi2> {
+// CHECK: %[[I8:.*]] = arith.trunci %[[IN]] : vector<8xi32> to vector<8xi8>
+// CHECK: %[[EVEN:.*]], %[[ODD:.*]] = vector.deinterleave %[[I8]] : vector<8xi8> -> vector<4xi8>
+// CHECK: vector.deinterleave %[[EVEN]] : vector<4xi8> -> vector<2xi8>
+// CHECK: vector.deinterleave %[[ODD]] : vector<4xi8> -> vector<2xi8>
+// CHECK-NOT: arith.trunci
+// CHECK: vector.bitcast {{.*}} : vector<2xi8> to vector<8xi2>
+ %0 = arith.trunci %a : vector<8xi32> to vector<8xi2>
+ return %0 : vector<8xi2>
+}
+
+// CHECK-LABEL: func.func @aligned_trunci_i8_to_i2_2d(
+func.func @aligned_trunci_i8_to_i2_2d(%a: vector<8x32xi8>) -> vector<8x32xi2> {
+// CHECK-SAME: %[[IN:.*]]: vector<8x32xi8>) -> vector<8x32xi2> {
+// CHECK-NOT: arith.trunci
+// CHECK: %[[EVEN:.*]], %[[ODD:.*]] = vector.deinterleave %[[IN]] : vector<8x32xi8> -> vector<8x16xi8>
+// CHECK: vector.deinterleave %[[EVEN]] : vector<8x16xi8> -> vector<8x8xi8>
+// CHECK: vector.deinterleave %[[ODD]] : vector<8x16xi8> -> vector<8x8xi8>
+// CHECK: vector.bitcast {{.*}} : vector<8x8xi8> to vector<8x32xi2>
+ %0 = arith.trunci %a : vector<8x32xi8> to vector<8x32xi2>
+ return %0 : vector<8x32xi2>
+}
+
+// CHECK-LABEL: func.func @unaligned_trunci_i8_to_i2(
+func.func @unaligned_trunci_i8_to_i2(%a: vector<6xi8>) -> vector<6xi2> {
+// CHECK-NOT: vector.deinterleave
+// CHECK: arith.trunci {{.*}} : vector<6xi8> to vector<6xi2>
+ %0 = arith.trunci %a : vector<6xi8> to vector<6xi2>
+ return %0 : vector<6xi2>
+}
+
///----------------------------------------------------------------------------------------
/// arith.extui
///
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/rewrite-narrow-types.mlir b/mlir/test/Integration/Dialect/Vector/CPU/rewrite-narrow-types.mlir
index 5f91ff5aacb1d4..bbe8d7ee749642 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/rewrite-narrow-types.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/rewrite-narrow-types.mlir
@@ -164,6 +164,27 @@ func.func @fext(%a: vector<5xi8>) {
return
}
+func.func @print_as_i1_8xi2(%v : vector<8xi2>) {
+ %bitsi16 = vector.bitcast %v : vector<8xi2> to vector<16xi1>
+ vector.print %bitsi16 : vector<16xi1>
+ return
+}
+
+func.func @ftrunc_i2(%a: vector<8xi8>) {
+ %0 = arith.trunci %a : vector<8xi8> to vector<8xi2>
+ func.call @print_as_i1_8xi2(%0) : (vector<8xi2>) -> ()
+ // CHECK: (
+ // CHECK-SAME: 1, 1,
+ // CHECK-SAME: 0, 1,
+ // CHECK-SAME: 1, 0,
+ // CHECK-SAME: 0, 0,
+ // CHECK-SAME: 1, 1,
+ // CHECK-SAME: 0, 1,
+ // CHECK-SAME: 1, 0,
+ // CHECK-SAME: 0, 0 )
+ return
+}
+
func.func @fcst_maskedload(%A: memref<?xi4>, %passthru: vector<6xi4>) -> vector<6xi4> {
%c0 = arith.constant 0: index
%mask = vector.constant_mask [3] : vector<6xi1>
@@ -194,6 +215,11 @@ func.func @entry() {
]> : vector<5xi8>
func.call @fext(%v4) : (vector<5xi8>) -> ()
+ %v5 = arith.constant dense<[
+ 0xef, 0xee, 0xed, 0xec, 0xeb, 0x02, 0x01, 0x00
+ ]> : vector<8xi8>
+ func.call @ftrunc_i2(%v5) : (vector<8xi8>) -> ()
+
// Set up memory.
%c0 = arith.constant 0: index
%c1 = arith.constant 1: index
>From ca4d552edb05740e785e86caa27721e0a4606bdb Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sun, 1 Nov 2026 20:09:47 +0800
Subject: [PATCH 2/2] Name the de-interleaved values as in the code and add an
i2 example
---
.../Transforms/VectorEmulateNarrowType.cpp | 29 ++++++++++++++++---
1 file changed, 25 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index c343ca92b7e0b5..b6e5e928fa7c44 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -1989,10 +1989,8 @@ static Value rewriteI8ToI2Trunc(PatternRewriter &rewriter, Location loc,
// the destination byte.
// in = [0,1,2,3,4,5,6,7],...
// evenOdd = [0,2,4,6],... [1,3,5,7],...
- // vec0 = [0,4],...
- // vec1 = [1,5],...
- // vec2 = [2,6],...
- // vec3 = [3,7],...
+ // vec02 = [0,4],... [2,6],...
+ // vec13 = [1,5],... [3,7],...
auto evenOdd = vector::DeinterleaveOp::create(rewriter, loc, srcValue);
auto vec02 = vector::DeinterleaveOp::create(rewriter, loc, evenOdd.getRes1());
auto vec13 = vector::DeinterleaveOp::create(rewriter, loc, evenOdd.getRes2());
@@ -2241,6 +2239,29 @@ struct RewriteAlignedSubByteIntExt : OpRewritePattern<ConversionOpType> {
/// %4 = arith.ori %2, %3 : vector<4xi8>
/// %5 = vector.bitcast %4 : vector<4xi8> to vector<8xi4>
///
+/// and, with four elements per byte,
+/// arith.trunci %in : vector<8xi32> to vector<8xi2>
+///
+/// is rewritten as:
+///
+/// %cst = arith.constant dense<3> : vector<2xi8>
+/// %cst_0 = arith.constant dense<2> : vector<2xi8>
+/// %cst_1 = arith.constant dense<4> : vector<2xi8>
+/// %cst_2 = arith.constant dense<6> : vector<2xi8>
+/// %0, %1 = vector.deinterleave %in : vector<8xi8> -> vector<4xi8>
+/// %2, %3 = vector.deinterleave %0 : vector<4xi8> -> vector<2xi8>
+/// %4, %5 = vector.deinterleave %1 : vector<4xi8> -> vector<2xi8>
+/// %6 = arith.andi %2, %cst : vector<2xi8>
+/// %7 = arith.andi %4, %cst : vector<2xi8>
+/// %8 = arith.shli %7, %cst_0 : vector<2xi8>
+/// %9 = arith.andi %3, %cst : vector<2xi8>
+/// %10 = arith.shli %9, %cst_1 : vector<2xi8>
+/// %11 = arith.shli %5, %cst_2 : vector<2xi8>
+/// %12 = arith.ori %6, %8 : vector<2xi8>
+/// %13 = arith.ori %12, %10 : vector<2xi8>
+/// %14 = arith.ori %13, %11 : vector<2xi8>
+/// %15 = vector.bitcast %14 : vector<2xi8> to vector<8xi2>
+///
struct RewriteAlignedSubByteIntTrunc : OpRewritePattern<arith::TruncIOp> {
using Base::Base;
More information about the Mlir-commits
mailing list