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

[Nishant Patel]

================
@@ -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,
----------------
nbpatel wrote:

ah yes, good point...I addressed the feedback, please take a look if this version looks good

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