[Mlir-commits] [mlir] [mlir][xegpu] Lower lane_data repack convert_layout to bitcast_shuffle (PR #210837)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 20 15:56:19 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
Extend the SgToLaneConvertLayout pattern to lower a convert_layout that only repacks lane_data between round-robin and contiguous form (keeping lane_layout and order unchanged) into xegpu.bitcast_shuffle. Each lane keeps the same elements and total bits along the repacked dimension, but their assignment to lanes changes, so the data is moved across lanes with
a bit-preserving bitcast_shuffle.
---
Patch is 22.10 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/210837.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td (+66)
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td (+5)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp (+29)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp (+150)
- (modified) mlir/test/Dialect/XeGPU/invalid.mlir (+14)
- (modified) mlir/test/Dialect/XeGPU/ops.mlir (+28)
- (modified) mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir (+102)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 7f8389a6acc47..0a6dc0b0240c5 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1491,6 +1491,72 @@ def XeGPU_TruncfOp
let hasVerifier = 1;
}
+def XeGPU_BitcastShuffleOp
+ : XeGPU_Op<"bitcast_shuffle", [Pure]> {
+ let summary = "Bit-preserving type conversion and re-distribution of a "
+ "subgroup's fragments across its lanes.";
+
+ let description = [{
+ `xegpu.bitcast_shuffle` is a lane-level operation. The `source` operand is
+ the fragment (a scalar or 1D vector) held by one lane, and the `result` is
+ the fragment that lane holds afterwards. It reinterprets the fragments as a
+ different element type, possibly with a different number of elements per
+ lane, and at the same time re-distributes the data across the lanes of the
+ subgroup.
+
+ Both `source` and `result` must be a scalar or a 1D vector of numerical
+ type, and they must have a different type. When they have the same number of
+ elements the operation degenerates into a plain bitcast. Otherwise the total
+ number of bits per lane must be equal on both sides.
+
+ For example, consider `vector<2xi16> -> i32` with a subgroup size of 8. Each
+ lane's `source` is a `vector<2xi16>` of two 16-bit half-words, so the
+ subgroup holds 8 * 2 = 16 half-words in total. Number these half-words 0..15
+ and denote a lane's `source` as `<a,b>`, so lane `i` starts with half-words
+ `<2*i, 2*i+1>`. The `result` `i32` holds two packed half-words, denoted
+ `a|b` (`a` in the low bits, `b` in the high bits):
+
+ ```
+ lane: 0 1 2 3 4 5 6 7
+ source: <0,1> <2,3> <4,5> <6,7> <8,9> <10,11> <12,13> <14,15>
+ result: 0|2 4|6 8|10 12|14 1|3 5|7 9|11 13|15
+ ```
+
+ The half-words are gathered one `source` element position at a time, in
+ lane order: first every lane's element 0 (half-words 0, 2, 4, ..., 14), then
+ every lane's element 1 (half-words 1, 3, 5, ..., 15). This sequence is then
+ handed back out two half-words per lane, in lane order, to form each lane's
+ `result`. In general each lane's two result half-words come from two
+ different lanes' `source` fragments -- that is the cross-lane shuffle.
+
+ The reverse pair `i32 -> vector<2xi16>` undoes this exactly, so bitcasting
+ to one type and back yields the original fragments.
+
+ The operation assumes a full subgroup (the subgroup size equals the maximum
+ subgroup size) and that all lanes execute the same dynamic instance of the
+ operation; otherwise the behavior is undefined.
+ }];
+
+ let arguments = (ins XeGPU_BitcastShuffleType:$source);
+ let results = (outs XeGPU_BitcastShuffleType:$result);
+
+ let extraClassDeclaration = [{
+ Type getSourceType() {
+ return getSource().getType();
+ }
+
+ Type getResultType() {
+ return getResult().getType();
+ }
+ }];
+
+ let assemblyFormat = [{
+ $source attr-dict `:` type($source) `->` type($result)
+ }];
+
+ let hasVerifier = 1;
+}
+
def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
AnchorLayoutInterface]> {
let summary = "It performs scaled mma computation";
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
index 0423303c23493..dfaea842ba4cc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
@@ -31,6 +31,11 @@ def XeGPU_VectorOrScalarType
: AnyTypeOf<[VectorOfRankAndType<[1,2,3,4,5,6,7,8], [XeGPU_ScalarType, Index]>, XeGPU_ScalarType]>;
def XeGPU_GatherScatterBaseAddrType
: AnyTypeOf<[MemRefRankOf<[XeGPU_ScalarType], [1]>, XeGPU_PointerType]>;
+// Operand/result type of xegpu.bitcast_shuffle: a numerical scalar or a 1D
+// vector of numerical type.
+def XeGPU_BitcastShuffleType
+ : AnyTypeOf<[FixedVectorOfRankAndType<[1], [XeGPU_ScalarType]>,
+ XeGPU_ScalarType]>;
// common base class for types in XeGPU dialect
class XeGPUTypeDef<string name, string typeMnemonic, list<Trait> traits = [],
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 2ffe883eb0d9a..24cd67720696f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -937,6 +937,35 @@ LogicalResult TruncfOp::verify() {
return success();
}
+//===----------------------------------------------------------------------===//
+// XeGPU_BitcastShuffleOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult BitcastShuffleOp::verify() {
+ Type srcType = getSourceType();
+ Type resType = getResultType();
+
+ // A bitcast_shuffle only makes sense between two different types; a
+ // same-type shuffle is a no-op.
+ if (srcType == resType)
+ return emitOpError("source and result must have different types.");
+
+ // Compute the total number of bits on both sides. The operation is
+ // bit-preserving, so the source and result must carry the same amount of
+ // data per lane.
+ auto sizeInBits = [](Type type) -> int64_t {
+ if (auto vecType = dyn_cast<VectorType>(type))
+ return vecType.getNumElements() * vecType.getElementTypeBitWidth();
+ return type.getIntOrFloatBitWidth();
+ };
+
+ if (sizeInBits(srcType) != sizeInBits(resType))
+ return emitOpError(
+ "source and result must have the same total size in bits.");
+
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// XeGPU_DpasMxOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 874487da10b30..509f9882decc8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1665,6 +1665,115 @@ shuffleDataAsLaneLayoutChange(ConversionPatternRewriter &rewriter, Location loc,
return res;
}
+/// Redistributes `src` for a `convert_layout` that repacks the `lane_data`
+/// along dimension `repackDim` between round-robin and contiguous, keeping
+/// `lane_layout` and `order` unchanged. Each lane keeps the same `k` elements
+/// and total bits along `repackDim`, but their assignment to lanes changes, so
+/// a `xegpu.bitcast_shuffle` moves the data across lanes: it converts each run
+/// of `k` elements between the round-robin form `vector<k x iElem>` (each lane
+/// owns `k` strided elements) and the contiguous form `vector<1 x iWide>` (each
+/// lane owns one packed `k * elemBits` value), with a `vector.bitcast` to and
+/// from the original element type on either side.
+///
+/// `inputData`/`targetData` are the `repackDim` `lane_data` of the input and
+/// target layouts; exactly one must be 1 (round-robin) and the other `k`
+/// (contiguous). Returns failure if that does not hold or the packed run would
+/// exceed the widest representable integer (64 bits).
+static FailureOr<Value>
+repackLaneData(ConversionPatternRewriter &rewriter, Location loc, Value src,
+ int64_t repackDim, int64_t inputData, int64_t targetData) {
+ auto srcTy = dyn_cast<VectorType>(src.getType());
+ if (!srcTy)
+ return failure();
+ int64_t rank = srcTy.getRank();
+ Type elemTy = srcTy.getElementType();
+ int64_t elemBits = elemTy.getIntOrFloatBitWidth();
+ int64_t k = srcTy.getShape()[repackDim];
+
+ bool roundRobinToContig = inputData == 1 && targetData == k;
+ bool contigToRoundRobin = inputData == k && targetData == 1;
+ if (!roundRobinToContig && !contigToRoundRobin)
+ return failure();
+
+ int64_t wideBits = k * elemBits;
+ if (wideBits > 64)
+ return failure();
+
+ MLIRContext *ctx = rewriter.getContext();
+ VectorType roundRobinIntTy = VectorType::get({k}, IntegerType::get(ctx, elemBits));
+ VectorType contigIntTy = VectorType::get({1}, IntegerType::get(ctx, wideBits));
+ VectorType runTy = VectorType::get({k}, elemTy);
+
+ // Repacks one run (vector<k x elemTy>) into the target form, returning a
+ // vector<k x elemTy>.
+ auto repackRun = [&](Value run) -> Value {
+ if (roundRobinToContig) {
+ Value asInt =
+ vector::BitCastOp::create(rewriter, loc, roundRobinIntTy, run);
+ Value shuffled =
+ xegpu::BitcastShuffleOp::create(rewriter, loc, contigIntTy, asInt);
+ return vector::BitCastOp::create(rewriter, loc, runTy, shuffled);
+ }
+ Value asWide = vector::BitCastOp::create(rewriter, loc, contigIntTy, run);
+ Value shuffled =
+ xegpu::BitcastShuffleOp::create(rewriter, loc, roundRobinIntTy, asWide);
+ return vector::BitCastOp::create(rewriter, loc, runTy, shuffled);
+ };
+
+ if (rank == 1)
+ return repackRun(src);
+
+ // When `repackDim` is innermost each run is a contiguous sub-vector and can
+ // be extracted and re-inserted as a whole.
+ if (repackDim == rank - 1) {
+ SmallVector<int64_t> outerShape(srcTy.getShape().drop_back());
+ int64_t numRuns = computeProduct(outerShape);
+ SmallVector<int64_t> outerStrides = computeStrides(outerShape);
+ Value result = arith::ConstantOp::create(rewriter, loc, srcTy,
+ rewriter.getZeroAttr(srcTy));
+ for (int64_t i = 0; i < numRuns; ++i) {
+ SmallVector<int64_t> pos = delinearize(i, outerStrides);
+ Value run = vector::ExtractOp::create(rewriter, loc, src, pos);
+ result = vector::InsertOp::create(rewriter, loc, repackRun(run), result,
+ pos);
+ }
+ return result;
+ }
+
+ // Otherwise each run is strided along `repackDim`, so it is gathered and
+ // scattered element by element for every combination of the kept indices.
+ SmallVector<int64_t> keptShape;
+ SmallVector<int64_t> keptDims;
+ for (int64_t d = 0; d < rank; ++d)
+ if (d != repackDim) {
+ keptShape.push_back(srcTy.getShape()[d]);
+ keptDims.push_back(d);
+ }
+ int64_t numRuns = computeProduct(keptShape);
+ SmallVector<int64_t> keptStrides = computeStrides(keptShape);
+ Value result = arith::ConstantOp::create(rewriter, loc, srcTy,
+ rewriter.getZeroAttr(srcTy));
+ for (int64_t i = 0; i < numRuns; ++i) {
+ SmallVector<int64_t> keptPos = delinearize(i, keptStrides);
+ SmallVector<int64_t> pos(rank, 0);
+ for (auto [dim, coord] : llvm::zip_equal(keptDims, keptPos))
+ pos[dim] = coord;
+ SmallVector<Value> runElems;
+ for (int64_t j = 0; j < k; ++j) {
+ pos[repackDim] = j;
+ runElems.push_back(vector::ExtractOp::create(rewriter, loc, src, pos));
+ }
+ Value repacked = repackRun(
+ vector::FromElementsOp::create(rewriter, loc, runTy, runElems));
+ for (int64_t j = 0; j < k; ++j) {
+ pos[repackDim] = j;
+ Value elem = vector::ExtractOp::create(rewriter, loc, repacked, j);
+ result = vector::InsertOp::create(rewriter, loc, elem, result, pos);
+ }
+ }
+ return result;
+}
+
/// Folds a subgroup-level ConvertLayout op with compatible lane layouts.
struct SgToLaneConvertLayout
: public OpConversionPattern<xegpu::ConvertLayoutOp> {
@@ -1721,6 +1830,47 @@ struct SgToLaneConvertLayout
}
}
+ // Handle a pure `lane_data` repack: `lane_layout` and `order` are unchanged
+ // and exactly one dimension's `lane_data` differs, switching between
+ // round-robin (lane_data 1) and contiguous (lane_data == run length). The
+ // number of elements per lane is unchanged, but their assignment to lanes
+ // is not, so the data is moved across lanes with `xegpu.bitcast_shuffle`.
+ //
+ // The changed dimension must be one of the two innermost ones: sg-to-lane
+ // distribution is 2D, with any higher (leading) dimensions being unit, so a
+ // repack of a 3rd-innermost-or-higher dimension is not supported.
+ if (inputLayout.getEffectiveOrderAsInt() ==
+ targetLayout.getEffectiveOrderAsInt() &&
+ inputLayout.getEffectiveLaneLayoutAsInt() ==
+ targetLayout.getEffectiveLaneLayoutAsInt()) {
+ // lane_layout and lane_data always share the same rank (enforced by the
+ // LayoutAttr verifier), so the two lane_data arrays have equal size here
+ // since their lane_layouts are equal.
+ auto laneData = inputLayout.getEffectiveLaneDataAsInt();
+ auto targetLaneData = targetLayout.getEffectiveLaneDataAsInt();
+ // Find the single dimension whose lane_data changed.
+ SmallVector<int64_t> changedDims;
+ for (auto [dim, data] : llvm::enumerate(llvm::zip_equal(
+ ArrayRef<int64_t>(laneData), ArrayRef<int64_t>(targetLaneData))))
+ if (std::get<0>(data) != std::get<1>(data))
+ changedDims.push_back(dim);
+
+ int64_t rank = laneData.size();
+ // Only a single dimension may change, and it must be one of the two
+ // innermost dimensions (>= rank - 2).
+ if (changedDims.size() == 1 && changedDims[0] >= rank - 2) {
+ int64_t repackDim = changedDims[0];
+ FailureOr<Value> res =
+ repackLaneData(rewriter, op.getLoc(), adaptor.getSource(),
+ repackDim, laneData[repackDim],
+ targetLaneData[repackDim]);
+ if (succeeded(res)) {
+ rewriter.replaceOp(op, *res);
+ return success();
+ }
+ }
+ }
+
return rewriter.notifyMatchFailure(
op, "lowering incompatible convert_layout not yet supported");
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index d5d4950fe7d7e..2ae6cacd0bf50 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -693,6 +693,20 @@ func.func @truncf_invalid_result_size(%a: vector<8x16xf16>) {
return
}
+// -----
+func.func @bitcast_shuffle_same_type(%a: vector<2xi16>) {
+ // expected-error at +1 {{op source and result must have different types}}
+ %1 = xegpu.bitcast_shuffle %a : vector<2xi16> -> vector<2xi16>
+ return
+}
+
+// -----
+func.func @bitcast_shuffle_size_mismatch(%a: i32) {
+ // expected-error at +1 {{op source and result must have the same total size in bits}}
+ %1 = xegpu.bitcast_shuffle %a : i32 -> vector<4xi16>
+ return
+}
+
// -----
func.func @dpas_mx_acc_result_type_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xbf16>) {
// expected-error at +1 {{Accumulator type must match result type.}}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index 6cffa3eec369b..3dbe570b3b1ad 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -688,6 +688,34 @@ gpu.func @truncf(%a: vector<8x16xf16>) {
gpu.return
}
+// CHECK-LABEL: gpu.func @bitcast_shuffle_scalar_to_vector_16
+gpu.func @bitcast_shuffle_scalar_to_vector_16(%a: i32) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : i32 -> vector<2xi16>
+ %1 = xegpu.bitcast_shuffle %a : i32 -> vector<2xi16>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @bitcast_shuffle_vector_16_to_scalar
+gpu.func @bitcast_shuffle_vector_16_to_scalar(%a: vector<2xi16>) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : vector<2xi16> -> i32
+ %1 = xegpu.bitcast_shuffle %a : vector<2xi16> -> i32
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @bitcast_shuffle_scalar_to_vector_8
+gpu.func @bitcast_shuffle_scalar_to_vector_8(%a: i32) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : i32 -> vector<4xi8>
+ %1 = xegpu.bitcast_shuffle %a : i32 -> vector<4xi8>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @bitcast_shuffle_vector_8_to_scalar
+gpu.func @bitcast_shuffle_vector_8_to_scalar(%a: vector<4xi8>) {
+ // CHECK: %{{.+}} = xegpu.bitcast_shuffle %{{.+}} : vector<4xi8> -> i32
+ %1 = xegpu.bitcast_shuffle %a : vector<4xi8> -> i32
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @dpas_mx
gpu.func @dpas_mx(%a : vector<8x32xf8E5M2>, %b: vector<32x16xf8E5M2>, %acc: vector<8x16xbf16>, %a_scale: vector<8x1xf8E8M0FNU>, %b_scale: vector<1x16xf8E8M0FNU>) {
// CHECK: %{{.+}} = xegpu.dpas_mx %{{.+}}, %{{.+}}, %{{.+}} scale_a = %{{.+}} scale_b = %{{.+}} : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index fe356e6af35c1..584d48d3a3a5e 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1062,6 +1062,108 @@ gpu.func @convert_layout_scalar() {
}
}
+// -----
+// A convert_layout that repacks the innermost lane_data from round-robin
+// ([1, 1]) to contiguous ([1, 4]) while keeping lane_layout unchanged lowers to
+// a per-row xegpu.bitcast_shuffle: each lane's 4 strided bf16 elements are
+// bitcast to i16, shuffled across lanes into one contiguous i64, and bitcast
+// back. The non-distributed outer dim (8 rows) is repacked independently.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_innermost_lane_data
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<8x64xbf16> to vector<8x4xbf16>
+// CHECK: %[[E0:.*]] = vector.extract %[[SRC]][0] : vector<4xbf16> from vector<8x4xbf16>
+// CHECK: %[[B0:.*]] = vector.bitcast %[[E0]] : vector<4xbf16> to vector<4xi16>
+// CHECK: %[[S0:.*]] = xegpu.bitcast_shuffle %[[B0]] : vector<4xi16> -> vector<1xi64>
+// CHECK: %{{.*}} = vector.bitcast %[[S0]] : vector<1xi64> to vector<4xbf16>
+// CHECK-COUNT-7: xegpu.bitcast_shuffle %{{.*}} : vector<4xi16> -> vector<1xi64>
+gpu.func @convert_layout_repack_innermost_lane_data() {
+ %src = "some_op"() : () -> vector<8x64xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>
+ }> : vector<8x64xbf16>
+ "some_use"(%cvt) : (vector<8x64xbf16>) -> ()
+ gpu.return
+}
+}
+
+// -----
+// The reverse repack (contiguous [1, 4] to round-robin [1, 1]) uses the
+// opposite bitcast_shuffle: each lane's contiguous i64 is shuffled back into 4
+// strided i16 elements.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_innermost_lane_data_reverse
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<8x64xbf16> to vector<8x4xbf16>
+// CHECK: %[[B0:.*]] = vector.bitcast %{{.*}} : vector<4xbf16> to vector<1xi64>
+// CHECK: %[[S0:.*]] = xegpu.bitcast_shuffle %[[B0]] : vector<1xi64> -> vector<4xi16>
+// CHECK: %{{.*}} = vector.bitcast %[[S0]] : vector<4xi16> to vector<4xbf16>
+gpu.func @convert_layout_repack_innermost_lane_data_reverse() {
+ %src = "some_op"() : () -> vector<8x64xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<8x64xbf16>
+ "some_use"(%cvt) : (vector<8x64xbf16>) -> ()
+ gpu.return
+}
+}
+
+// -----
+// A 1D convert_layout repack is a single bitcast_shuffle with no outer loop.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_innermost_lane_data_1d
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<64xbf16> to vector<4xbf16>
+// CHECK: %[[B0:.*]] = vector.bitcast %[[SRC]] : vector<4xbf16> to vector<4xi16>
+// CHECK: %[[S0:.*]] = xegpu.bitcast_shuffle %[[B0]] : vector<4xi16> -> vector<1xi64>
+// CHECK: %{{.*}} = vector.bitcast %[[S0]] : vector<1xi64> to vector<4xbf16>
+gpu.func @convert_layout_repack_innermost_lane_data_1d() {
+ %src = "some_op"() : () -> vector<64xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lan...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/210837
More information about the Mlir-commits
mailing list