[Mlir-commits] [mlir] [mlir][XeGPU] Add an SLM round-trip fallback for sg-to-lane convert_layout (PR #227884)

Jianhui Li llvmlistbot at llvm.org
Wed Sep 30 14:48:19 PDT 2026


https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/227884

This PR adds a general fallback lowering for xegpu.convert_layout in the sg-to-lane distribution pass, which round-trips the value through shared local memory.

The existing lane-level lowerings are all special cases: the layouts fold into each other, or they differ in a way a xegpu.lane_shuffle or a gpu.shuffle can express. Anything else failed to legalize and stopped the pipeline — for instance a conversion that only moves the lane_data of a dimension distributed over part of the subgroup:

```mlir
xegpu.convert_layout %src
  <{input_layout  = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 4]>,
    target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>}>
  : vector<1x2x32xbf16>
```

Here dim 2 is distributed over 8 of the 16 lanes and dim 1 over the other factor of 2, so the subgroup splits into two groups that each hold a run of their own. A single xegpu.lane_shuffle round-robins with a stride of the whole subgroup size and cannot express that.

The new pattern works for any pair of layouts which can both distribute the value: each lane writes its fragment to SLM at the coordinates input_layout assigns it, then reads it back from the ones target_layout assigns. It mirrors what XeGPUWgToSgDistribute already does for a workgroup-level convert_layout whose sg_layout changes, and is registered at a lower pattern benefit than every specialized lowering, so it only runs once they have all declined.

assisted-by-claude

>From 58d77046557189c81bea19971b8eab67ed635ff1 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 30 Sep 2026 21:18:03 +0000
Subject: [PATCH] [mlir][XeGPU] Add an SLM round-trip fallback for sg-to-lane
 convert_layout

The lane-level `convert_layout` lowerings are all special cases: the layouts
fold into each other, or they differ in a way a `xegpu.lane_shuffle` or a
`gpu.shuffle` can express. Anything else failed to legalize and stopped the
pipeline, for instance a conversion that only moves the `lane_data` of a
dimension distributed over part of the subgroup:

  xegpu.convert_layout %src
    <{input_layout  = #xegpu.layout<lane_layout = [1, 2, 8],
                                    lane_data = [1, 1, 4]>,
      target_layout = #xegpu.layout<lane_layout = [1, 2, 8],
                                    lane_data = [1, 1, 1]>}>
    : vector<1x2x32xbf16>

Add a last-resort lowering that works for any pair of layouts which can both
distribute the value: write each lane's fragment to shared local memory at the
coordinates `input_layout` assigns it, then read it back from the ones
`target_layout` assigns. This mirrors what `XeGPUWgToSgDistribute` already does
for a workgroup-level `convert_layout` whose `sg_layout` changes. It is
registered at a lower pattern benefit than every specialized lowering, so it
only runs once they have all declined.

Two details differ from the workgroup-level version. The op is subgroup level
but every subgroup of the workgroup executes it on its own data, so the scratch
holds one tile per subgroup, stacked along the outermost dimension and addressed
through a `gpu.subgroup_id` offset; the subgroup count comes from the kernel's
`known_block_size` and the pattern declines when it is absent. And the exchange
is between the lanes of one subgroup, which is a single hardware thread, so
ordering the write before the read with `xegpu.fence` is enough and no
`gpu.barrier` is needed.

A lane's fragment is one `lane_data`-sized piece per distribution unit, so the
transfer is one `store_matrix`/`load_matrix` pair per unit. The existing
`load_matrix`/`store_matrix` lane distribution asserts on layouts with more than
one unit, which is why the pieces are sliced out here rather than handed to those
patterns.

Deriving the subgroup count from `known_block_size` was already open coded in
XeGPUPropagateLayout, so it moves to `xegpu::getNumSubgroupsFromBlockSize` and
both callers share it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h     |   7 +
 .../XeGPU/Transforms/XeGPUPropagateLayout.cpp |  16 +-
 .../Transforms/XeGPUSgToLaneDistribute.cpp    | 202 ++++++++++++++++++
 mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp   |  18 ++
 .../XeGPU/sg-to-lane-distribute-unit.mlir     |  51 +++++
 5 files changed, 282 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index d54d159fe2120..4db69bda113dc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -104,6 +104,13 @@ Value createVectorWithShapeFromValues(OpBuilder &builder, Location loc,
 /// if no GPU module parent or XeVM target attribute exists.
 std::optional<std::string> getChipStr(Operation *op);
 
+/// Returns the number of subgroups the kernel enclosing `op` runs, derived from
+/// the `known_block_size` of its parent `gpu.func`. Fails when there is no such
+/// parent, when the attribute is absent, when a block dimension is not a
+/// positive power of two, or when the block does not cover a whole subgroup.
+FailureOr<int64_t> getNumSubgroupsFromBlockSize(Operation *op,
+                                                int64_t subgroupSize);
+
 /// Generates element-wise addition ops of two arrays with same length.
 SmallVector<OpFoldResult> addElementwise(OpBuilder &builder, Location loc,
                                          ArrayRef<OpFoldResult> lhs,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index f05193504116b..08512f67a87e1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -486,18 +486,10 @@ FailureOr<int64_t> LayoutInfoPropagation::getNumSgOrFail(
       return llvm::product_of(sgLayout);
   }
   // Otherwise fall back to the kernel's known_block_size.
-  if (auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>()) {
-    std::optional<ArrayRef<int32_t>> knownBlockSize =
-        gpuFunc.getKnownBlockSize();
-    if (knownBlockSize) {
-      bool isPowerOf2Block = llvm::all_of(*knownBlockSize, [](int32_t dim) {
-        return dim > 0 && llvm::isPowerOf2_32(dim);
-      });
-      int64_t numSg = llvm::product_of(*knownBlockSize) / sgSize;
-      if (isPowerOf2Block && numSg > 0)
-        return numSg;
-    }
-  }
+  if (FailureOr<int64_t> numSg =
+          xegpu::getNumSubgroupsFromBlockSize(op, sgSize);
+      succeeded(numSg))
+    return *numSg;
   // Only subgroup mode needs the count; elsewhere a missing one is benign.
   if (layoutKind == xegpu::LayoutKind::Subgroup) {
     markFailure(op, "Unable to determine the number of subgroups for the "
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 06f769cefb4fd..539175be768ac 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1939,6 +1939,203 @@ struct SgToLaneConvertLayout
   }
 };
 
+/// The offsets and shape of one `lane_data`-sized piece of a lane's fragment.
+/// `computeDistributedCoords` returns one coordinate per distribution unit,
+/// enumerated row major over `shape / (lane_layout * lane_data)`, and unit `u`
+/// owns the piece of the fragment at `delinearize(u) * lane_data`.
+struct LaneFragmentPiece {
+  SmallVector<int64_t> offsets;
+  SmallVector<int64_t> sizes;
+};
+
+static LaneFragmentPiece getLaneFragmentPiece(int64_t unit,
+                                              ArrayRef<int64_t> shape,
+                                              ArrayRef<int64_t> laneLayout,
+                                              ArrayRef<int64_t> laneData) {
+  int64_t rank = shape.size();
+  SmallVector<int64_t> units(rank);
+  for (int64_t d = 0; d < rank; ++d)
+    units[d] = shape[d] / std::min(shape[d], laneLayout[d] * laneData[d]);
+  SmallVector<int64_t> unitCoords = delinearize(unit, computeStrides(units));
+  LaneFragmentPiece piece;
+  for (int64_t d = 0; d < rank; ++d) {
+    piece.offsets.push_back(unitCoords[d] * laneData[d]);
+    piece.sizes.push_back(laneData[d]);
+  }
+  return piece;
+}
+
+/// Last-resort lowering for a `convert_layout` that none of the patterns above
+/// handle: round-trip the value through shared local memory, writing each
+/// lane's fragment to the coordinates `input_layout` gives it and reading it
+/// back from the ones `target_layout` gives it. Any pair of layouts that can
+/// both distribute the value works, at the cost of an SLM write and read.
+///
+/// The op is subgroup level, but every subgroup of the workgroup executes it on
+/// its own data, so the scratch holds one tile per subgroup and each subgroup
+/// addresses its own through a `gpu.subgroup_id` offset on the outermost
+/// dimension. The subgroup count comes from the kernel's `known_block_size`.
+///
+/// A lane's fragment is one `lane_data`-sized piece per distribution unit, so
+/// the transfer is one `store_matrix`/`load_matrix` per unit; layouts with a
+/// single unit, which is the common case, give a single pair.
+struct SgToLaneConvertLayoutViaSLM
+    : public OpConversionPattern<xegpu::ConvertLayoutOp> {
+  using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(xegpu::ConvertLayoutOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
+    xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
+    if (!inputLayout || !targetLayout || !inputLayout.isForSubgroup() ||
+        !targetLayout.isForSubgroup())
+      return rewriter.notifyMatchFailure(op, "both layouts must be lane level");
+
+    auto valueTy = dyn_cast<VectorType>(op.getResult().getType());
+    if (!valueTy)
+      return rewriter.notifyMatchFailure(op, "value type must be a vector");
+    Type elemTy = valueTy.getElementType();
+    // The scratch is sized in bytes, so sub-byte elements would not get a
+    // well-defined address.
+    if (!elemTy.isIntOrFloat() || elemTy.getIntOrFloatBitWidth() % 8 != 0)
+      return rewriter.notifyMatchFailure(
+          op, "element type must be a whole number of bytes");
+
+    ArrayRef<int64_t> sgShape = valueTy.getShape();
+    int64_t rank = sgShape.size();
+    if (rank != inputLayout.getRank() || rank != targetLayout.getRank())
+      return rewriter.notifyMatchFailure(
+          op, "both layouts must have the rank of the value");
+
+    FailureOr<VectorType> distInputTy =
+        xegpu::getDistVecTypeBasedOnLaneLayout(inputLayout, valueTy);
+    FailureOr<VectorType> distTargetTy =
+        xegpu::getDistVecTypeBasedOnLaneLayout(targetLayout, valueTy);
+    if (failed(distInputTy) || failed(distTargetTy))
+      return rewriter.notifyMatchFailure(
+          op, "value type must be distributable by both layouts");
+
+    const auto *uArch =
+        xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+    if (!uArch)
+      return rewriter.notifyMatchFailure(
+          op, "target attribute is required to determine the subgroup size");
+    FailureOr<int64_t> numSubgroups =
+        xegpu::getNumSubgroupsFromBlockSize(op, uArch->getSubgroupSize());
+    if (failed(numSubgroups))
+      return rewriter.notifyMatchFailure(
+          op, "the scratch holds one tile per subgroup, so @known_block_size "
+              "must be attached to the kernel, with power-of-two dimensions "
+              "covering at least one subgroup");
+
+    Location loc = op.getLoc();
+
+    // One tile per subgroup, stacked along the outermost dimension.
+    SmallVector<int64_t> slmShape(sgShape);
+    slmShape[0] *= *numSubgroups;
+    int64_t slmBytes =
+        computeProduct(slmShape) * elemTy.getIntOrFloatBitWidth() / 8;
+    auto slmTy = MemRefType::get({slmBytes}, rewriter.getI8Type(), {}, 3);
+    Value slm = memref::AllocaOp::create(rewriter, loc, slmTy);
+    Value memDesc = xegpu::CreateMemDescOp::create(
+        rewriter, loc,
+        xegpu::MemDescType::get(rewriter.getContext(), slmShape, elemTy,
+                                nullptr),
+        slm);
+
+    // This subgroup's tile starts at `subgroup_id` tiles into the scratch.
+    Value sgId = gpu::SubgroupIdOp::create(rewriter, loc,
+                                           rewriter.getIndexType(), nullptr);
+    SmallVector<Value> base(rank);
+    base[0] = arith::MulIOp::create(
+        rewriter, loc, sgId,
+        arith::ConstantIndexOp::create(rewriter, loc, sgShape[0]));
+    for (int64_t d = 1; d < rank; ++d)
+      base[d] = arith::ConstantIndexOp::create(rewriter, loc, 0);
+
+    Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+                                         /*upperBound=*/mlir::IntegerAttr());
+    auto dynamicOffsets = rewriter.getDenseI64ArrayAttr(
+        SmallVector<int64_t>(rank, ShapedType::kDynamic));
+
+    // Adds this subgroup's base to one distribution unit's coordinates.
+    auto withBase = [&](ArrayRef<Value> coords) {
+      SmallVector<Value> offsets;
+      for (auto [coord, b] : llvm::zip_equal(coords, base))
+        offsets.push_back(arith::AddIOp::create(rewriter, loc, coord, b));
+      return offsets;
+    };
+
+    // Write phase: every lane stores the pieces `input_layout` gives it.
+    auto storeCoords =
+        inputLayout.computeDistributedCoords(rewriter, loc, laneId, sgShape);
+    if (failed(storeCoords))
+      return rewriter.notifyMatchFailure(
+          op, "failed to compute the input_layout coordinates");
+    SmallVector<int64_t> inputLaneLayout =
+        inputLayout.getEffectiveLaneLayoutAsInt();
+    SmallVector<int64_t> inputLaneData =
+        inputLayout.getEffectiveLaneDataAsInt();
+    Value fragment =
+        castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getSource()),
+                    *distInputTy);
+    for (auto [unit, coords] : llvm::enumerate(*storeCoords)) {
+      LaneFragmentPiece piece =
+          getLaneFragmentPiece(unit, sgShape, inputLaneLayout, inputLaneData);
+      Value data = fragment;
+      if (piece.sizes != SmallVector<int64_t>(distInputTy->getShape()))
+        data = vector::ExtractStridedSliceOp::create(
+            rewriter, loc, fragment, piece.offsets, piece.sizes,
+            SmallVector<int64_t>(rank, 1));
+      xegpu::StoreMatrixOp::create(rewriter, loc, TypeRange{}, data, memDesc,
+                                   ValueRange(withBase(coords)), dynamicOffsets,
+                                   nullptr, xegpu::DistributeLayoutAttr{});
+    }
+
+    // The scratch is only exchanged between the lanes of one subgroup, which is
+    // a single hardware thread, so ordering the write before the read is enough
+    // and no workgroup barrier is needed.
+    xegpu::FenceOp::create(rewriter, loc, xegpu::MemorySpace::SLM,
+                           xegpu::FenceScope::Workgroup);
+
+    // Read phase: every lane loads the pieces `target_layout` gives it.
+    auto loadCoords =
+        targetLayout.computeDistributedCoords(rewriter, loc, laneId, sgShape);
+    if (failed(loadCoords))
+      return rewriter.notifyMatchFailure(
+          op, "failed to compute the target_layout coordinates");
+    SmallVector<int64_t> targetLaneLayout =
+        targetLayout.getEffectiveLaneLayoutAsInt();
+    SmallVector<int64_t> targetLaneData =
+        targetLayout.getEffectiveLaneDataAsInt();
+    Value result;
+    if (loadCoords->size() > 1)
+      result = arith::ConstantOp::create(rewriter, loc, *distTargetTy,
+                                         rewriter.getZeroAttr(*distTargetTy));
+    for (auto [unit, coords] : llvm::enumerate(*loadCoords)) {
+      LaneFragmentPiece piece =
+          getLaneFragmentPiece(unit, sgShape, targetLaneLayout, targetLaneData);
+      auto pieceTy = VectorType::get(piece.sizes, elemTy);
+      Value loaded = xegpu::LoadMatrixOp::create(
+          rewriter, loc, pieceTy, memDesc, ValueRange(withBase(coords)),
+          dynamicOffsets, nullptr, xegpu::DistributeLayoutAttr{});
+      if (!result) {
+        result = loaded;
+        break;
+      }
+      result = vector::InsertStridedSliceOp::create(
+          rewriter, loc, loaded, result, piece.offsets,
+          SmallVector<int64_t>(rank, 1));
+    }
+
+    rewriter.replaceOp(op, castValueTo(rewriter,
+                                       cast<TypedValue<VectorType>>(result),
+                                       *distTargetTy));
+    return success();
+  }
+};
+
 /// `getEffectiveLaneDataAsInt` is empty when `lane_data` is unset, so the unit
 /// check also rejects layouts that are not lane level.
 static bool hasDefaultOrderAndUnitLaneData(xegpu::DistributeLayoutAttr layout) {
@@ -2687,4 +2884,9 @@ void xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
            SgToLaneConvertLayoutPartialBroadcastExtractShuffle,
            SgToLaneConvertLayoutDeinterleaveSelect>(typeConverter,
                                                     patterns.getContext());
+  // The SLM round-trip handles any pair of layouts but pays for a memory
+  // round-trip, so it only runs when every pattern above has declined.
+  patterns.add<SgToLaneConvertLayoutViaSLM>(typeConverter,
+                                            patterns.getContext(),
+                                            /*benefit=*/PatternBenefit(0));
 }
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index bf50fd3b3af7e..d1d4f044c8ed0 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -490,6 +490,24 @@ std::optional<std::string> xegpu::getChipStr(Operation *op) {
   return std::nullopt;
 }
 
+FailureOr<int64_t> xegpu::getNumSubgroupsFromBlockSize(Operation *op,
+                                                       int64_t subgroupSize) {
+  auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>();
+  if (!gpuFunc)
+    return failure();
+  std::optional<ArrayRef<int32_t>> blockSize = gpuFunc.getKnownBlockSize();
+  if (!blockSize)
+    return failure();
+  if (!llvm::all_of(*blockSize, [](int32_t dim) {
+        return dim > 0 && llvm::isPowerOf2_32(dim);
+      }))
+    return failure();
+  int64_t numSubgroups = llvm::product_of(*blockSize) / subgroupSize;
+  if (numSubgroups < 1)
+    return failure();
+  return numSubgroups;
+}
+
 /// Generates element-wise addition ops of two arrays with same length.
 SmallVector<OpFoldResult> xegpu::addElementwise(OpBuilder &builder,
                                                 Location loc,
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 e69fabbce0dc3..cb5872a0973fb 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1283,6 +1283,57 @@ gpu.func @convert_layout_repack_innermost_lane_data_3d() {
 }
 }
 
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL:     gpu.func @convert_layout_via_slm_lane_layout
+// CHECK-NOT:         xegpu.convert_layout
+// CHECK:             %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<16x16xf32> to vector<16x1xf32>
+// CHECK:             %[[SLM:.*]] = memref.alloca() : memref<4096xi8, 3>
+// CHECK:             %[[MD:.*]] = xegpu.create_mem_desc %[[SLM]] : memref<4096xi8, 3> -> !xegpu.mem_desc<64x16xf32>
+// CHECK:             %[[SGID:.*]] = gpu.subgroup_id : index
+// CHECK:             %[[ROWS:.*]] = arith.constant 16 : index
+// CHECK:             %[[BASE:.*]] = arith.muli %[[SGID]], %[[ROWS]] : index
+// CHECK:             gpu.lane_id
+// CHECK-COUNT-16:    xegpu.store_matrix %{{.*}}, %[[MD]]{{\[}}%{{.*}}, %{{.*}}] : vector<1x1xf32>, !xegpu.mem_desc<64x16xf32>, index, index
+// CHECK:             xegpu.fence memory_kind = slm, fence_scope = workgroup
+// CHECK-COUNT-16:    xegpu.load_matrix %[[MD]]{{\[}}%{{.*}}, %{{.*}}] : !xegpu.mem_desc<64x16xf32>, index, index -> vector<1x1xf32>
+// CHECK:             builtin.unrealized_conversion_cast %{{.*}} : vector<1x16xf32> to vector<16x16xf32>
+gpu.func @convert_layout_via_slm_lane_layout() kernel attributes {known_block_size = array<i32: 64, 1, 1>} {
+  %src = "test.some_op"() : () -> vector<16x16xf32>
+  %cvt = xegpu.convert_layout %src
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+      target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+    }> : vector<16x16xf32>
+  "test.some_use"(%cvt) : (vector<16x16xf32>) -> ()
+  gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL:     gpu.func @convert_layout_via_slm_lane_data
+// CHECK-NOT:         xegpu.convert_layout
+// CHECK:             %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<1x2x32xbf16> to vector<1x1x4xbf16>
+// CHECK:             %[[SLM:.*]] = memref.alloca() : memref<512xi8, 3>
+// CHECK:             %[[MD:.*]] = xegpu.create_mem_desc %[[SLM]] : memref<512xi8, 3> -> !xegpu.mem_desc<4x2x32xbf16>
+// CHECK:             gpu.subgroup_id : index
+// CHECK:             xegpu.store_matrix %[[SRC]], %[[MD]]{{\[}}%{{.*}}, %{{.*}}, %{{.*}}] : vector<1x1x4xbf16>, !xegpu.mem_desc<4x2x32xbf16>, index, index, index
+// CHECK:             xegpu.fence memory_kind = slm, fence_scope = workgroup
+// CHECK-COUNT-4:     xegpu.load_matrix %[[MD]]{{\[}}%{{.*}}, %{{.*}}, %{{.*}}] : !xegpu.mem_desc<4x2x32xbf16>, index, index, index -> vector<1x1x1xbf16>
+// CHECK:             builtin.unrealized_conversion_cast %{{.*}} : vector<1x1x4xbf16> to vector<1x2x32xbf16>
+gpu.func @convert_layout_via_slm_lane_data() kernel attributes {known_block_size = array<i32: 64, 1, 1>} {
+  %src = "test.some_op"() : () -> vector<1x2x32xbf16>
+  %cvt = xegpu.convert_layout %src
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 4]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>
+    }> : vector<1x2x32xbf16>
+  "test.some_use"(%cvt) : (vector<1x2x32xbf16>) -> ()
+  gpu.return
+}
+}
+
 // -----
 // load_matrix and store_matrix with coordinate computation (offsets [0,0])
 gpu.module @xevm_module {



More information about the Mlir-commits mailing list