[Mlir-commits] [mlir] [MLIR][XeGPU] Add non-splat constant distribution in SgToLane pass (PR #205575)

[Nishant Patel]

https://github.com/nbpatel created https://github.com/llvm/llvm-project/pull/205575

None

>From 0b3f5844be4ec0ca4d1d0fac641c4b299536d1ed Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 22 Jun 2026 21:15:54 +0000
Subject: [PATCH] Add non-splat constant distribution for SgToLane

---
 .../Transforms/XeGPUSgToLaneDistribute.cpp    | 68 ++++++++++++++++---
 .../XeGPU/sg-to-lane-distribute-unit.mlir     | 20 ++++++
 2 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 75a87f84b3da8..25fcd6054e569 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -355,6 +355,13 @@ struct SgToLaneElementWise : public ConversionPattern {
 
 /// Distributes a subgroup-level arith ConstantOp to lane-level arith
 /// ConstantOp.
+///
+/// Splat constants are distributed by simply rebuilding the splat with the
+/// lane-local vector type. Non-splat dense constants are distributed as :
+/// `computeDistributedCoords` yields the coordinates each lane owns, each
+/// element is extracted from the full (subgroup-level) constant, and the
+/// per-lane elements are assembled into the distributed vector with
+/// `vector.from_elements` (or `vector.broadcast` for a single element).
 struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
   using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;
 
@@ -365,11 +372,11 @@ struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
     if (!resultType)
       return failure();
 
-    // Only handle dense vector constants
-    auto dense = dyn_cast<SplatElementsAttr>(op.getValue());
-    if (!dense)
+    // Only handle dense vector constants.
+    auto denseAttr = dyn_cast<DenseElementsAttr>(op.getValue());
+    if (!denseAttr)
       return rewriter.notifyMatchFailure(
-          op, "only dense splat vector constants are supported");
+          op, "only dense vector constants are supported");
 
     xegpu::DistributeLayoutAttr layout =
         xegpu::getTemporaryLayout(llvm::cast<OpResult>(op.getResult()));
@@ -385,12 +392,55 @@ struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
           op, "unable to compute lane vector type from the layout");
 
     VectorType newResultType = laneShapeOrFailure.value();
-    auto sclarValue = dense.getSplatValue<Attribute>();
-    auto newDenseAttr = DenseElementsAttr::get(newResultType, sclarValue);
+    Location loc = op.getLoc();
 
-    auto newOp = arith::ConstantOp::create(rewriter, op.getLoc(), newResultType,
-                                           newDenseAttr);
-    rewriter.replaceOp(op, newOp.getResult());
+    // Splat constants: every lane gets the same value, so just rebuild the
+    // splat with the distributed type.
+    if (denseAttr.isSplat()) {
+      auto scalarValue = denseAttr.getSplatValue<Attribute>();
+      auto newDenseAttr = DenseElementsAttr::get(newResultType, scalarValue);
+      auto newOp =
+          arith::ConstantOp::create(rewriter, loc, newResultType, newDenseAttr);
+      rewriter.replaceOp(op, newOp.getResult());
+      return success();
+    }
+
+    // Non-splat constants: each lane extracts the elements it owns from the
+    // full constant using the distributed coordinates from the layout.
+    auto fullConst =
+        arith::ConstantOp::create(rewriter, loc, resultType, denseAttr);
+
+    Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+                                         /*upperBound=*/mlir::IntegerAttr());
+    auto maybeCoordsVec = layout.computeDistributedCoords(
+        rewriter, loc, laneId, resultType.getShape());
+    if (failed(maybeCoordsVec))
+      return rewriter.notifyMatchFailure(
+          op, "failed to compute distributed coordinates from layout");
+
+    SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
+    int64_t numElements = newResultType.getNumElements();
+    assert(static_cast<int64_t>(coordsVec.size()) == numElements &&
+           "number of coordinate sets must match number of distributed "
+           "elements");
+
+    SmallVector<Value> elements;
+    for (auto &coords : coordsVec) {
+      SmallVector<OpFoldResult> mixedPos = getAsOpFoldResult(coords);
+      elements.push_back(vector::ExtractOp::create(
+          rewriter, loc, fullConst.getResult(), mixedPos));
+    }
+
+    // Assemble the distributed vector.
+    Value result;
+    if (numElements == 1) {
+      result = vector::BroadcastOp::create(rewriter, loc, newResultType,
+                                           elements[0]);
+    } else {
+      result = vector::FromElementsOp::create(rewriter, loc, newResultType,
+                                              elements);
+    }
+    rewriter.replaceOp(op, result);
     return success();
   }
 };
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 fcc2da3d5005b..61f4cc5ceb675 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -160,6 +160,26 @@ gpu.func @arith_constant() {
   gpu.return
 }
 
+// A non-splat constant is distributed by extracting the element each lane owns
+// from the full constant (using the layout's distributed coordinates).
+// CHECK-LABEL: gpu.func @arith_constant_non_splat
+// CHECK: %[[CST:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xindex>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ELEM:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<16xindex>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ELEM]] : index to vector<1xindex>
+// CHECK: gpu.return
+gpu.func @arith_constant_non_splat() {
+  %0 = arith.constant
+    {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}
+    dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xindex>
+  %cl0 = xegpu.convert_layout %0
+    <{
+      input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,
+      target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>
+    }> : vector<16xindex>
+  gpu.return
+}
+
 // CHECK-LABEL: gpu.func @prefetch_nd
 // CHECK: %[[C0:.*]] = arith.constant 0 : index
 // CHECK: xegpu.prefetch_nd %{{.*}}[%[[C0]], %[[C0]]] : !xegpu.tensor_desc<16x16xf16>