[Mlir-commits] [mlir] [MLIR][XeGPU] Hoist i4 extensions above lane shuffles (PR #234760)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 22 13:50:45 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Sang Ik Lee (silee2)
<details>
<summary>Changes</summary>
Hoist the i4-to-i8 portion of integer extensions through single-use lane-shuffle and vector data-movement chains before XeGPU-to-XeVM conversion. Legalize narrow casts into byte-level operations in a separate preprocessing stage so the existing bitcast-shuffle lowering can handle the widened fragments without repacking.
Cover signed and unsigned extensions, supported fragment widths, layout-distribution shuffles, and unsupported consumers. Scope preprocessing to casts and generated operations to preserve existing conversion diagnostics.
---
Patch is 21.41 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/234760.diff
5 Files Affected:
- (modified) mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt (+1)
- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+112)
- (modified) mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir (+57)
- (added) mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_ext.mlir (+133)
- (added) mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_layout.mlir (+60)
``````````diff
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
index dd9edc43a16570..3590fd65b246e1 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
+++ b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
@@ -17,6 +17,7 @@ add_mlir_conversion_library(MLIRXeGPUToXeVM
MLIRLLVMDialect
MLIRXeVMDialect
MLIRVectorDialect
+ MLIRVectorTransforms
MLIRArithDialect
MLIRIndexDialect
MLIRSCFDialect
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index b2a81f01887bb9..1c7e1d24397390 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -21,11 +21,15 @@
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
+#include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"
@@ -1559,6 +1563,78 @@ class TruncfToXeVMPattern : public OpConversionPattern<arith::TruncFOp> {
}
};
+// Hoist the i4-to-i8 part of an extension through a single-use data-movement
+// chain and lane shuffle. Widening all the way to the extension's result type
+// could exceed the bitcast-shuffle lowering's 64-bit fragment limit.
+// Signed widening preserves the eventual signed interpretation of each nibble;
+// unsigned widening preserves its unsigned interpretation.
+template <typename ExtOp>
+class HoistI4ExtensionAboveLaneShuffle final : public OpRewritePattern<ExtOp> {
+public:
+ using OpRewritePattern<ExtOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ExtOp op,
+ PatternRewriter &rewriter) const override {
+ auto srcTy = dyn_cast<VectorType>(op.getIn().getType());
+ auto dstTy = dyn_cast<VectorType>(op.getType());
+ if (!srcTy || !srcTy.getElementType().isInteger(4) || !dstTy ||
+ dstTy.getElementTypeBitWidth() < 8)
+ return failure();
+
+ SmallVector<Operation *> movers;
+ Value source = op.getIn();
+ while (Operation *def = source.getDefiningOp()) {
+ if (!source.hasOneUse())
+ return failure();
+ if (isa<xegpu::LaneShuffleOp>(def))
+ break;
+ if (!isa<vector::ShapeCastOp, vector::ExtractStridedSliceOp,
+ vector::ExtractOp, vector::BroadcastOp, vector::TransposeOp>(def))
+ return failure();
+ // Only vector results are supported; scalar extraction needs separate
+ // narrow-type legalization.
+ if (!isa<VectorType>(def->getOperand(0).getType()) ||
+ !isa<VectorType>(def->getResult(0).getType()))
+ return failure();
+ movers.push_back(def);
+ source = def->getOperand(0);
+ }
+ auto shuffle = source.getDefiningOp<xegpu::LaneShuffleOp>();
+ if (!shuffle || !source.hasOneUse())
+ return failure();
+ VectorType fragmentTy = shuffle.getSourceType();
+ int64_t numElements = fragmentTy.getNumElements();
+ if (!fragmentTy.getElementType().isInteger(4) || numElements < 2 ||
+ numElements > 8 || !llvm::isPowerOf2_64(numElements))
+ return failure();
+
+ rewriter.setInsertionPoint(shuffle);
+ VectorType wideTy = fragmentTy.clone(rewriter.getI8Type());
+ Value wide =
+ ExtOp::create(rewriter, op.getLoc(), wideTy, shuffle.getSource());
+ Value result = xegpu::LaneShuffleOp::create(
+ rewriter, shuffle.getLoc(), wideTy, wide, shuffle.getMode());
+ for (Operation *mover : llvm::reverse(movers)) {
+ rewriter.setInsertionPoint(mover);
+ IRMapping mapping;
+ mapping.map(mover->getOperand(0), result);
+ Operation *clone = rewriter.clone(*mover, mapping);
+ clone->getResult(0).setType(
+ cast<VectorType>(mover->getResult(0).getType())
+ .clone(rewriter.getI8Type()));
+ result = clone->getResult(0);
+ }
+ if (result.getType() == dstTy)
+ rewriter.replaceOp(op, result);
+ else
+ rewriter.modifyOpInPlace(op, [&] { op.getInMutable().assign(result); });
+ for (Operation *mover : movers)
+ rewriter.eraseOp(mover);
+ rewriter.eraseOp(shuffle);
+ return success();
+ }
+};
+
// Lowers `xegpu.lane_shuffle` to `xevm.bitcast_shuffle`.
//
// `xevm.bitcast_shuffle` concatenates the components of its operand across the
@@ -1636,6 +1712,42 @@ struct ConvertXeGPUToXeVMPass
void runOnOperation() override {
MLIRContext *context = &getContext();
+ // Restrict greedy rewriting to casts and newly generated operations so
+ // unrelated dead XeGPU ops still receive conversion diagnostics.
+ GreedyRewriteConfig config;
+ config.setStrictness(GreedyRewriteStrictness::ExistingAndNewOps);
+ // Preserve extension consumers until hoisting has finished: narrow-type
+ // legalization would otherwise obscure the chains this peephole matches.
+ {
+ RewritePatternSet patterns(context);
+ patterns.add<HoistI4ExtensionAboveLaneShuffle<arith::ExtUIOp>,
+ HoistI4ExtensionAboveLaneShuffle<arith::ExtSIOp>>(context);
+ SmallVector<Operation *> extensions;
+ getOperation()->walk([&](Operation *op) {
+ if (isa<arith::ExtUIOp, arith::ExtSIOp>(op))
+ extensions.push_back(op);
+ });
+ if (failed(
+ applyOpPatternsGreedily(extensions, std::move(patterns), config)))
+ return signalPassFailure();
+ }
+ // Expand the new i4-to-i8 casts into byte-level masks and shifts before
+ // lowering. LLVM extensions on i4 vectors are not legal in the backend.
+ {
+ RewritePatternSet patterns(context);
+ vector::populateVectorNarrowTypeRewritePatterns(patterns);
+ vector::populateVectorInterleaveToShufflePatterns(patterns);
+ vector::populateVectorDeinterleaveToShufflePatterns(patterns);
+ SmallVector<Operation *> casts;
+ getOperation()->walk([&](Operation *op) {
+ if (isa<arith::ExtUIOp, arith::ExtSIOp, arith::TruncIOp,
+ vector::BitCastOp>(op))
+ casts.push_back(op);
+ });
+ if (failed(applyOpPatternsGreedily(casts, std::move(patterns), config)))
+ return signalPassFailure();
+ }
+
// XeVM type converter is based on LLVM type converter with the
// following customizations.
// First, type conversion rules are added for xegpu custom types,
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index 74e8964d752f70..63ee2170de024e 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -76,6 +76,63 @@ gpu.module @test_kernel {
// -----
+// A non-extension consumer cannot hoist widening above an i4 shuffle.
+
+gpu.module @test_kernel {
+ gpu.func @lane_shuffle_i4_non_extension(%a: vector<4xi4>) -> vector<4xi4> {
+ // expected-error at +1 {{failed to legalize operation 'xegpu.lane_shuffle' that was explicitly marked illegal}}
+ %0 = xegpu.lane_shuffle %a unpack : vector<4xi4>
+ gpu.return %0 : vector<4xi4>
+ }
+}
+
+// -----
+
+// Mixed signed and unsigned extensions cannot share one widened shuffle.
+
+gpu.module @test_kernel {
+ gpu.func @lane_shuffle_i4_mixed_uses(%a: vector<2xi8>) -> (vector<4xi32>, vector<4xi32>) {
+ %nibbles = vector.bitcast %a : vector<2xi8> to vector<4xi4>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.lane_shuffle' that was explicitly marked illegal}}
+ %0 = xegpu.lane_shuffle %nibbles pack : vector<4xi4>
+ %unsigned = arith.extui %0 : vector<4xi4> to vector<4xi32>
+ %signed = arith.extsi %0 : vector<4xi4> to vector<4xi32>
+ gpu.return %unsigned, %signed : vector<4xi32>, vector<4xi32>
+ }
+}
+
+// -----
+
+// The single-use restriction also applies to intermediate data movers.
+
+gpu.module @test_kernel {
+ gpu.func @lane_shuffle_i4_shape_cast_multiple_uses(%a: vector<2xi8>) -> (vector<1x4xi32>, vector<1x4xi4>) {
+ %nibbles = vector.bitcast %a : vector<2xi8> to vector<4xi4>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.lane_shuffle' that was explicitly marked illegal}}
+ %0 = xegpu.lane_shuffle %nibbles unpack : vector<4xi4>
+ %shaped = vector.shape_cast %0 : vector<4xi4> to vector<1x4xi4>
+ %extended = arith.extui %shaped : vector<1x4xi4> to vector<1x4xi32>
+ gpu.return %extended, %shaped : vector<1x4xi32>, vector<1x4xi4>
+ }
+}
+
+// -----
+
+// Even with a single extension consumer, widening sixteen i4 elements to i8
+// would exceed the byte shuffle's 64-bit fragment limit.
+
+gpu.module @test_kernel {
+ gpu.func @lane_shuffle_i4_ext_too_wide(%a: vector<8xi8>) -> vector<16xi32> {
+ %nibbles = vector.bitcast %a : vector<8xi8> to vector<16xi4>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.lane_shuffle' that was explicitly marked illegal}}
+ %0 = xegpu.lane_shuffle %nibbles pack : vector<16xi4>
+ %extended = arith.extui %0 : vector<16xi4> to vector<16xi32>
+ gpu.return %extended : vector<16xi32>
+ }
+}
+
+// -----
+
// Verify that a xegpu.lane_shuffle fragment wider than 64 bits is rejected.
gpu.module @test_kernel {
diff --git a/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_ext.mlir b/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_ext.mlir
new file mode 100644
index 00000000000000..7a7a5fcee2e3ea
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_ext.mlir
@@ -0,0 +1,133 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm -split-input-file %s | FileCheck %s --implicit-check-not='arith.{{(extui|extsi|trunci)}} {{.*}}xi4>' --implicit-check-not=arith.ori --implicit-check-not=xegpu.lane_shuffle
+
+// Hoist only the i4-to-i8 part of the extension before the shuffle. Extending
+// to i32 first would exceed the shuffle's 64-bit fragment limit. Starting with
+// packed bytes also checks that narrow-type legalization unpacks the nibbles
+// without repacking them after the shuffle.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @unpack_shape_cast_extui
+// CHECK-SAME: %[[SRC:.*]]: vector<2xi8>
+gpu.func @unpack_shape_cast_extui(%a: vector<2xi8>) -> vector<1x4xi32> {
+ // CHECK: %[[MASK:.*]] = arith.constant dense<15> : vector<2xi8>
+ // CHECK: %[[LOW:.*]] = arith.andi %[[SRC]], %[[MASK]] : vector<2xi8>
+ // CHECK: %[[SHIFT:.*]] = arith.constant dense<4> : vector<2xi8>
+ // CHECK: %[[HIGH:.*]] = arith.shrui %[[SRC]], %[[SHIFT]] : vector<2xi8>
+ // CHECK: %[[BYTES:.*]] = vector.shuffle %[[LOW]], %[[HIGH]] [0, 2, 1, 3] : vector<2xi8>, vector<2xi8>
+ // CHECK: %[[PACKED:.*]] = llvm.bitcast %[[BYTES]] : vector<4xi8> to i32
+ // CHECK: %[[SHUFFLED:.*]] = xevm.bitcast_shuffle %[[PACKED]] : (i32) -> vector<4xi8>
+ // CHECK: %[[SHAPED:.*]] = vector.shape_cast %[[SHUFFLED]] : vector<4xi8> to vector<1x4xi8>
+ // CHECK: %[[RES:.*]] = arith.extui %[[SHAPED]] : vector<1x4xi8> to vector<1x4xi32>
+ // CHECK: gpu.return %[[RES]] : vector<1x4xi32>
+ %nibbles = vector.bitcast %a : vector<2xi8> to vector<4xi4>
+ %shuffled = xegpu.lane_shuffle %nibbles unpack : vector<4xi4>
+ %shaped = vector.shape_cast %shuffled : vector<4xi4> to vector<1x4xi4>
+ %res = arith.extui %shaped : vector<1x4xi4> to vector<1x4xi32>
+ gpu.return %res : vector<1x4xi32>
+}
+}
+
+// -----
+
+// Signed widening must sign-extend each nibble before packing the byte vector,
+// then preserve the signed extension from i8 to i32 after the shuffle.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @pack_extsi
+// CHECK-SAME: %[[SRC:.*]]: vector<2xi8>
+gpu.func @pack_extsi(%a: vector<2xi8>) -> vector<4xi32> {
+ // CHECK: %[[SHIFT:.*]] = arith.constant dense<4> : vector<2xi8>
+ // CHECK: %[[LEFT:.*]] = arith.shli %[[SRC]], %[[SHIFT]] : vector<2xi8>
+ // CHECK: %[[LOW:.*]] = arith.shrsi %[[LEFT]], %{{.*}} : vector<2xi8>
+ // CHECK: %[[HIGH:.*]] = arith.shrsi %[[SRC]], %{{.*}} : vector<2xi8>
+ // CHECK: %[[BYTES:.*]] = vector.shuffle %[[LOW]], %[[HIGH]] [0, 2, 1, 3] : vector<2xi8>, vector<2xi8>
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %[[BYTES]] : (vector<4xi8>) -> i32
+ // CHECK: %[[SHUFFLED:.*]] = llvm.bitcast %[[PACKED]] : i32 to vector<4xi8>
+ // CHECK: %[[RES:.*]] = arith.extsi %[[SHUFFLED]] : vector<4xi8> to vector<4xi32>
+ // CHECK: gpu.return %[[RES]] : vector<4xi32>
+ %nibbles = vector.bitcast %a : vector<2xi8> to vector<4xi4>
+ %shuffled = xegpu.lane_shuffle %nibbles pack : vector<4xi4>
+ %res = arith.extsi %shuffled : vector<4xi4> to vector<4xi32>
+ gpu.return %res : vector<4xi32>
+}
+}
+
+// -----
+
+// An extension whose destination is already i8 disappears entirely.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @unpack_extui_i8
+gpu.func @unpack_extui_i8(%a: vector<2xi8>) -> vector<4xi8> {
+ // CHECK: %[[PACKED:.*]] = llvm.bitcast %{{.*}} : vector<4xi8> to i32
+ // CHECK: %[[RES:.*]] = xevm.bitcast_shuffle %[[PACKED]] : (i32) -> vector<4xi8>
+ // CHECK-NEXT: gpu.return %[[RES]] : vector<4xi8>
+ %nibbles = vector.bitcast %a : vector<2xi8> to vector<4xi4>
+ %shuffled = xegpu.lane_shuffle %nibbles unpack : vector<4xi4>
+ %res = arith.extui %shuffled : vector<4xi4> to vector<4xi8>
+ gpu.return %res : vector<4xi8>
+}
+}
+
+// -----
+
+// The smallest supported fragment widens to 16 bits. Signed extension to i8
+// needs no remaining cast either.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @pack_two_extsi_i8
+gpu.func @pack_two_extsi_i8(%a: vector<1xi8>) -> vector<2xi8> {
+ // CHECK: arith.shrsi {{.*}} : vector<1xi8>
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %{{.*}} : (vector<2xi8>) -> i16
+ // CHECK: %[[RES:.*]] = llvm.bitcast %[[PACKED]] : i16 to vector<2xi8>
+ // CHECK-NEXT: gpu.return %[[RES]] : vector<2xi8>
+ %nibbles = vector.bitcast %a : vector<1xi8> to vector<2xi4>
+ %shuffled = xegpu.lane_shuffle %nibbles pack : vector<2xi4>
+ %res = arith.extsi %shuffled : vector<2xi4> to vector<2xi8>
+ gpu.return %res : vector<2xi8>
+}
+}
+
+// -----
+
+// Eight elements reach the byte shuffle's 64-bit limit. Hoist through a slice
+// followed by a shape cast, keeping the slice and final extension after it.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @unpack_eight_slice_shape_extui
+gpu.func @unpack_eight_slice_shape_extui(%a: vector<4xi8>) -> vector<1x4xi32> {
+ // CHECK: %[[PACKED:.*]] = llvm.bitcast %{{.*}} : vector<8xi8> to i64
+ // CHECK: %[[SHUFFLED:.*]] = xevm.bitcast_shuffle %[[PACKED]] : (i64) -> vector<8xi8>
+ // CHECK: %[[SLICE:.*]] = vector.extract_strided_slice %[[SHUFFLED]] offsets = [2], sizes = [4], strides = [1] : vector<8xi8> to vector<4xi8>
+ // CHECK: %[[SHAPED:.*]] = vector.shape_cast %[[SLICE]] : vector<4xi8> to vector<1x4xi8>
+ // CHECK: %[[RES:.*]] = arith.extui %[[SHAPED]] : vector<1x4xi8> to vector<1x4xi32>
+ // CHECK: gpu.return %[[RES]] : vector<1x4xi32>
+ %nibbles = vector.bitcast %a : vector<4xi8> to vector<8xi4>
+ %shuffled = xegpu.lane_shuffle %nibbles unpack : vector<8xi4>
+ %slice = vector.extract_strided_slice %shuffled offsets = [2], sizes = [4], strides = [1] : vector<8xi4> to vector<4xi4>
+ %shaped = vector.shape_cast %slice : vector<4xi4> to vector<1x4xi4>
+ %res = arith.extui %shaped : vector<1x4xi4> to vector<1x4xi32>
+ gpu.return %res : vector<1x4xi32>
+}
+}
+
+// -----
+
+// Exercise the remaining supported movers in a single-use chain: transpose,
+// vector-valued extract, and broadcast. None may force nibble repacking.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @pack_transpose_extract_broadcast_extsi
+gpu.func @pack_transpose_extract_broadcast_extsi(%a: vector<2xi8>) -> vector<3x2xi16> {
+ // CHECK: %[[PACKED:.*]] = xevm.bitcast_shuffle %{{.*}} : (vector<4xi8>) -> i32
+ // CHECK: %[[SHUFFLED:.*]] = llvm.bitcast %[[PACKED]] : i32 to vector<4xi8>
+ // CHECK: %[[SHAPED:.*]] = vector.shape_cast %[[SHUFFLED]] : vector<4xi8> to vector<2x2xi8>
+ // CHECK: %[[TRANSPOSED:.*]] = vector.transpose %[[SHAPED]], [1, 0] : vector<2x2xi8> to vector<2x2xi8>
+ // CHECK: %[[ROW:.*]] = vector.extract %[[TRANSPOSED]][1] : vector<2xi8> from vector<2x2xi8>
+ // CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[ROW]] : vector<2xi8> to vector<3x2xi8>
+ // CHECK: %[[RES:.*]] = arith.extsi %[[BROADCAST]] : vector<3x2xi8> to vector<3x2xi16>
+ // CHECK: gpu.return %[[RES]] : vector<3x2xi16>
+ %nibbles = vector.bitcast %a : vector<2xi8> to vector<4xi4>
+ %shuffled = xegpu.lane_shuffle %nibbles pack : vector<4xi4>
+ %shaped = vector.shape_cast %shuffled : vector<4xi4> to vector<2x2xi4>
+ %transposed = vector.transpose %shaped, [1, 0] : vector<2x2xi4> to vector<2x2xi4>
+ %row = vector.extract %transposed[1] : vector<2xi4> from vector<2x2xi4>
+ %broadcast = vector.broadcast %row : vector<2xi4> to vector<3x2xi4>
+ %res = arith.extsi %broadcast : vector<3x2xi4> to vector<3x2xi16>
+ gpu.return %res : vector<3x2xi16>
+}
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_layout.mlir b/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_layout.mlir
new file mode 100644
index 00000000000000..36cf0a19a57abe
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/lane_shuffle_i4_layout.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-opt --test-xegpu-sg-to-lane-distribute %s | FileCheck %s --check-prefix=DIST
+// RUN: mlir-opt --test-xegpu-sg-to-lane-distribute --convert-xegpu-to-xevm %s | FileCheck %s --implicit-check-not='arith.{{(extui|extsi|trunci)}} {{.*}}xi4>' --implicit-check-not=arith.ori --implicit-check-not=xegpu.lane_shuffle
+
+// A contiguous four-nibble lane fragment is unpacked to round-robin order,
+// then extended for a vector store. Slice one row from the 8x64 tile so the
+// layout conversion produces one shuffle, without unsupported vector inserts.
+// Exercise the shuffle produced by layout distribution with a real sub-byte
+// block load rather than a fragment argument. The test distribution pass does
+// not distribute vector.store, so it materializes its input with a cast.
+gpu.module @test [#xevm.target<chip = "cri">] {
+ // DIST-LABEL: gpu.func @load_convert_layout_extui
+ // DIST: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<32xi4>
+ // DIST: %[[TILE:.*]] = vector.shape_cast %[[LOAD]] : vector<32xi4> to vector<8x4xi4>
+ // DIST: %[[ROW:.*]] = vector.extract_strided_slice %[[TILE]] offsets = [0, 0], sizes = [1, 4], strides = [1, 1] : vector<8x4xi4> to vector<1x4xi4>
+ // DIST: %[[FLAT:.*]] = vector.shape_cast %[[ROW]] : vector<1x4xi4> to vector<4xi4>
+ // DIST: %[[SHUFFLED:.*]] = xegpu.lane_shuffle %[[FLAT]] unpack : vector<4xi4>
+ // DIST: %[[SHAPED:.*]] = vector.shape_cast %[[SHUFFLED]] : vector<4xi4> to vector<1x4xi4>
+ // DIST: %[[EXT:.*]] = arith.extui %[[SHAPED]] : vector<1x4xi4> to vector<1x4xi32>
+ // DIST: %[[STORE:.*]] = builtin.unrealized_conversion_cast %[[EXT]] : vector<1x4xi32> to vector<1x64xi32>
+ // DIST: vector.store %[[STORE]]
+ // CHECK-LABEL: gpu.func @load_convert_layout_extui
+ // CHECK-SAME: %[[SRC:.*]]: memref<8x64xi4>, %[[DST:.*]]: memref<8x64xi32>
+ gpu.func @load_convert_layout_extui(%src: memref<8x64xi4>, %dst: memref<8x64xi32>) {
+ // CHECK: %[[LOAD:.*]] = xevm.blockload2d
+ // CHECK-SAME: elem_size_in_bits = 16 : i32
+ // CHECK-SAME: tile_height = 8 : i32, tile_width = 16 : i32
+ // CHECK-SAME: -> vector<8xi16>
+ // CHECK: %[[NIBBLES:.*]] = vector.bitcast %[[LOAD]] : vector<8xi16> to vector<32xi4>
+ // CHECK: %[[TILE:.*]] = vector.shape_cast %[[NIBBLES]] : vector<32xi4> to vector<8x4xi4>
+ // CHECK: %[[ROW:.*]] = vector.extract_strided_slice %[[TILE]] offsets = [0, 0], sizes = [1, 4], strides = [1, 1] : vector<8x4xi4> to vector<1x4xi4>
+ // CHECK: %[[FLAT:.*]] = vector.shape_cast %[[ROW]] : vector<1x4xi4> to vector<4xi4>
+ // CHECK: %[[BYTES:.*]] = vector.bitcast %[[FLAT]] : vector<4xi4> to vector<2xi8>
+ // CHECK: %[[LOW:.*]] = arith.andi %[[BYTES]], %{{.*}} : vector<2xi8>
+ // CHECK: %[[HIGH:.*]] = arith.shrui %[[BYTES]], %{{.*}} : vector<2xi8>
+ // CHECK: %[[WIDE:.*]] = vector.shuffle %[[LOW]], %[[HIGH]] [0, 2, 1, 3] : vector<2xi8>, vector<2xi8>
+ // CHECK: %[[PACKED:.*]] = llvm.bitcast %[[WIDE]] : vector<4xi8> to i32
+ // CHECK: %[[SHU...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/234760
More information about the Mlir-commits
mailing list