[Mlir-commits] [mlir] [MLIR][XeGPU] Fold lane_data repack convert_layout (PR #204016)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Sang Ik Lee (silee2)

<details>
<summary>Changes</summary>

into extract_strided_slice consumers

Fold a subgroup-level convert_layout that merely repacks lane_data along the non-distributed outer dimension (from [N, 1] to [1, 1] with order = [1, 0]), keeping lane_layout unchanged, when its result is consumed by exactly N vector.extract_strided_slice ops. After lane distribution both layouts yield the same per-lane vector (getDistributedVectorType ignores lane_data) and the N slices recover the original blocking, so the convert is redundant and folds to its source.

Add a lit test covering the fold in sg-to-lane-distribute-unit.mlir.

---
Full diff: https://github.com/llvm/llvm-project/pull/204016.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp (+39) 
- (modified) mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir (+55) 


``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 75a87f84b3da8..12a9b71220987 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1638,6 +1638,45 @@ struct SgToLaneConvertLayout
       }
     }
 
+    // Special case: fold a convert_layout that merely repacks `lane_data`
+    // along the non-distributed outer dimension - from `[N, 1]` to `[1, 1]`
+    // with `order = [1, 0]`, keeping `lane_layout` unchanged - when its result
+    // is consumed by exactly `N` `vector.extract_strided_slice` ops. After lane
+    // distribution both layouts yield the same per-lane vector
+    // (`getDistributedVectorType` ignores `lane_data`), and the `N` slices
+    // recover the original blocking, so the convert is redundant and folds to
+    // its source.
+    if (inputLayout.getRank() == 2 && targetLayout.getRank() == 2) {
+      auto laneLayout = inputLayout.getEffectiveLaneLayoutAsInt();
+      auto targetLaneLayout = targetLayout.getEffectiveLaneLayoutAsInt();
+      auto laneData = inputLayout.getEffectiveLaneDataAsInt();
+      auto targetLaneData = targetLayout.getEffectiveLaneDataAsInt();
+      auto targetOrder = targetLayout.getEffectiveOrderAsInt();
+      if (laneLayout.size() == 2 && targetLaneLayout.size() == 2 &&
+          laneData.size() == 2 && targetLaneData.size() == 2 &&
+          laneLayout == targetLaneLayout && laneLayout[0] == 1 &&
+          laneData[0] > 1 && laneData[1] == 1 &&
+          targetLaneData == SmallVector<int64_t>({1, 1}) &&
+          targetOrder == SmallVector<int64_t>({1, 0})) {
+        // The block factor `N` is the outer (non-distributed) lane_data, and
+        // it must match the number of extract_strided_slice consumers.
+        int64_t blockFactor = laneData[0];
+        int64_t numSlices = 0;
+        bool allSlices = true;
+        for (Operation *user : op.getResult().getUsers()) {
+          ++numSlices;
+          if (!isa<vector::ExtractStridedSliceOp>(user)) {
+            allSlices = false;
+            break;
+          }
+        }
+        if (allSlices && numSlices == blockFactor) {
+          rewriter.replaceOp(op, adaptor.getSource());
+          return success();
+        }
+      }
+    }
+
     return rewriter.notifyMatchFailure(
         op, "lowering incompatible convert_layout not yet supported");
   }
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 56db482f033d3..fcc2da3d5005b 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -794,6 +794,61 @@ gpu.func @vector_extract_strided_slice_partial_offsets() {
   gpu.return
 }
 
+// A convert_layout that only repacks lane_data along the non-distributed outer
+// dimension (from [4, 1] to [1, 1] with order = [1, 0]), keeping lane_layout
+// unchanged, folds to its source when consumed by exactly 4 (== outer
+// lane_data) extract_strided_slice ops. After distribution both layouts yield
+// the same per-lane vector, so the convert is redundant and the slices operate
+// directly on the distributed source.
+// CHECK-LABEL: gpu.func @convert_layout_repack_lane_data
+// CHECK-NOT:     xegpu.convert_layout
+// CHECK:         %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<32x16xi8> to vector<32x1xi8>
+// CHECK:         vector.extract_strided_slice %[[SRC]] {offsets = [0, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+// CHECK:         vector.extract_strided_slice %[[SRC]] {offsets = [8, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+// CHECK:         vector.extract_strided_slice %[[SRC]] {offsets = [16, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+// CHECK:         vector.extract_strided_slice %[[SRC]] {offsets = [24, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+gpu.func @convert_layout_repack_lane_data() {
+  %src = "some_op"() : () -> vector<32x16xi8>
+  %cvt = xegpu.convert_layout %src
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+    }> : vector<32x16xi8>
+  %s0 = vector.extract_strided_slice %cvt
+    {offsets = [0, 0], sizes = [8, 16], strides = [1, 1]}
+    : vector<32x16xi8> to vector<8x16xi8>
+  %s1 = vector.extract_strided_slice %cvt
+    {offsets = [8, 0], sizes = [8, 16], strides = [1, 1]}
+    : vector<32x16xi8> to vector<8x16xi8>
+  %s2 = vector.extract_strided_slice %cvt
+    {offsets = [16, 0], sizes = [8, 16], strides = [1, 1]}
+    : vector<32x16xi8> to vector<8x16xi8>
+  %s3 = vector.extract_strided_slice %cvt
+    {offsets = [24, 0], sizes = [8, 16], strides = [1, 1]}
+    : vector<32x16xi8> to vector<8x16xi8>
+  %a0 = xegpu.convert_layout %s0
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+    }> : vector<8x16xi8>
+  %a1 = xegpu.convert_layout %s1
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+    }> : vector<8x16xi8>
+  %a2 = xegpu.convert_layout %s2
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+    }> : vector<8x16xi8>
+  %a3 = xegpu.convert_layout %s3
+    <{
+      input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+      target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+    }> : vector<8x16xi8>
+  gpu.return
+}
+
 // CHECK-LABEL: gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted
 // CHECK: %[[ISS:.*]] = vector.insert_strided_slice %{{.*}}, %{{.*}} {offsets = [24, 0], strides = [1, 1]} : vector<16x1xf32> into vector<64x1xf32>
 gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted() {

``````````

</details>


https://github.com/llvm/llvm-project/pull/204016