[Mlir-commits] [mlir] [MLIR][XeGPU] Make shape cast respect full replication layout semantics (PR #229395)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir

Author: Artem Kroviakov (akroviakov)

<details>
<summary>Changes</summary>

A non-distributed dimension is also the one that does full replication across lanes, not only `layout[d] == 1`. 
Lanes on the replicated dimension share the same offset in tiling.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+18-2) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+15) 


``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 031d971cebe6c..2ea24aac2c375 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2851,11 +2851,21 @@ xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
       consumerLayout.getEffectiveLaneLayoutAsInt();
   SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
 
+  // A dim is not distributed when its lanes cover a single coordinate.
+  // TODO: partially replicated dims (1 < shape_dim / lane_data < lane_layout)
+  // are still treated as fully distributed.
+  auto isUndistributed = [&](int64_t dim, ArrayRef<int64_t> lanes,
+                             ArrayRef<int64_t> data) {
+    return lanes[dim] == 1 || resShape[dim] / data[dim] <= 1;
+  };
+
   xegpu::DistributeLayoutAttr resLayout = [&]() -> DistributeLayoutAttr {
+    SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
     DistributeLayoutAttr layout = consumerLayout;
     for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
       size_t i = 0;
-      while (i < dimGroup.size() && laneLayout[dimGroup[i]] == 1)
+      while (i < dimGroup.size() &&
+             isUndistributed(dimGroup[i], laneLayout, laneData))
         ++i;
 
       for (++i; i < dimGroup.size(); ++i) {
@@ -2876,6 +2886,12 @@ xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
   SmallVector<int64_t> resInstData = resLayout.getEffectiveInstDataAsInt();
   SmallVector<int64_t> resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
   SmallVector<int64_t> resLaneData = resLayout.getEffectiveLaneDataAsInt();
+  // `isContiguousTiling` assumes lanes have unique positions.
+  // Treat replicated dims the same way as single-lane dims.
+  SmallVector<int64_t> resSteppingLanes(resLaneLayout);
+  for (int64_t dim = 0, rank = resSteppingLanes.size(); dim < rank; ++dim)
+    if (isUndistributed(dim, resLaneLayout, resLaneData))
+      resSteppingLanes[dim] = 1;
 
   for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
     // One lane's elements must be a single contiguous run of the collapsed dim.
@@ -2887,7 +2903,7 @@ xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
     int64_t collapsedLaneData = 1;
     for (int64_t dim : dimGroup)
       collapsedLaneData *= resLaneData[dim];
-    if (!isContiguousTiling(dimGroup, resShape, resLaneLayout, resLaneData,
+    if (!isContiguousTiling(dimGroup, resShape, resSteppingLanes, resLaneData,
                             /*expectedStride=*/collapsedLaneData))
       return nullptr;
 
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 30e02cc42095c..a0b261f12402d 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -772,3 +772,18 @@ func.func @shape_cast_split_fills_strided_inner_dim(%arg0: memref<16x1024xbf16>)
   return
 }
 }
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @shape_cast_insert_unit_dim_replicated_lanes(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 4, 4], lane_data = [1, 1, 1]>, dims = [1]>} dense<0.000000e+00> : vector<8x32xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 4, 4], lane_data = [1, 1, 1]>} : vector<8x32xf16> to vector<8x1x32xf16>
+func.func @shape_cast_insert_unit_dim_replicated_lanes() {
+  %cst = arith.constant dense<0.000000e+00> : vector<8x32xf16>
+  %0 = vector.shape_cast %cst : vector<8x32xf16> to vector<8x1x32xf16>
+  %1 = xegpu.convert_layout %0
+     <{target_layout = #xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 4, 4], lane_data = [1, 1, 1]>}>
+     : vector<8x1x32xf16>
+  return
+}
+}

``````````

</details>


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