[Mlir-commits] [mlir] [MLIR][XeGPU] Optimize create_mask sg-to-lane lowering (PR #221198)

[Artem Kroviakov]

================
@@ -933,36 +913,90 @@ struct SgToLaneCreateMask : public OpConversionPattern<OpType> {
       return rewriter.notifyMatchFailure(
           op, "failed to compute distributed coordinates from layout");
 
-    SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
+    SmallVector<SmallVector<Value>> laneDataCoords = maybeCoordsVec.value();
+    SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
+    ArrayRef<int64_t> distShape = distType.getShape();
+    int64_t rank = distType.getRank();
     int64_t numElements = distType.getNumElements();
-    assert(static_cast<int64_t>(coordsVec.size()) == numElements &&
-           "number of coordinate sets must match number of distributed "
-           "elements");
-
-    // For each element, compare all coordinates against bounds.
-    Value trueVal =
-        arith::ConstantIntOp::create(rewriter, loc, /*value=*/1, /*width=*/1);
-    SmallVector<Value> maskBits;
-    for (auto &coords : coordsVec) {
-      Value inBounds = trueVal;
-      for (size_t i = 0; i < coords.size(); ++i) {
-        Value cmp = arith::CmpIOp::create(
-            rewriter, loc, arith::CmpIPredicate::slt, coords[i], origBounds[i]);
-        inBounds = arith::AndIOp::create(rewriter, loc, inBounds, cmp);
+
+    if (static_cast<int64_t>(laneData.size()) != rank ||
+        !computeShapeRatio(distShape, laneData))
+      return rewriter.notifyMatchFailure(
+          op, "lane_data does not tile the distributed vector");
+
+    SmallVector<int64_t> resultStrides = computeStrides(distShape);
+    assert(static_cast<int64_t>(laneDataCoords.size()) *
+                   computeProduct(laneData) ==
+               numElements &&
+           "number of coordinate sets must match number of lane_data blocks");
+
+    // Static offsets to be applied to a dynamic lane's dist units coordinates.
+    SmallVector<SmallVector<int64_t>> staticLaneDataOffsetsOrigShape =
+        layout.computeStaticDistributedCoords(/*linearId=*/0, origShape);
+    if (staticLaneDataOffsetsOrigShape.empty() ||
+        staticLaneDataOffsetsOrigShape.size() != laneDataCoords.size())
+      return rewriter.notifyMatchFailure(
+          op, "static and dynamic coordinates disagree on the distribution "
+              "unit count");
+
+    SmallVector<int64_t> unitTile(rank, 1);
+    int64_t linearLaneDataIdx = 0;
+    // Layout and shape information is static, compute static offset of lane's
+    // elements in the source dimensions.
+    SmallVector<SmallVector<int64_t>> staticElemOffsetInSrcDims(
+        rank, SmallVector<int64_t>(numElements));
+    // For each lane_data block in the overall lane's shape
+    for (SmallVector<int64_t> laneDataOffsetInResult :
+         StaticTileOffsetRange(distShape, laneData)) {
+      ArrayRef<int64_t> staticLaneDataOffsetInSource =
+          staticLaneDataOffsetsOrigShape[linearLaneDataIdx++];
+      // For each element in the lane_data block
+      for (SmallVector<int64_t> elementOffsetInLaneData :
+           StaticTileOffsetRange(laneData, unitTile)) {
+        // For each dim of indexing space
+        SmallVector<int64_t> elementOffsetInResult(rank);
+        for (int64_t d = 0; d < rank; d++)
+          elementOffsetInResult[d] =
+              laneDataOffsetInResult[d] + elementOffsetInLaneData[d];
+        int64_t elemLinearizedIdx =
+            linearize(elementOffsetInResult, resultStrides);
+        for (int64_t d = 0; d < rank; d++)
+          staticElemOffsetInSrcDims[d][elemLinearizedIdx] =
----------------
akroviakov wrote:

> While elementLinearizedIdx seems already incorporate rank d info, but still need put to staticElemOffsetInSrcDims[d]?

The two indices live in different spaces. elemLinearizedIdx is a linear position in the distributed shape.

Renamed and added a detailed pattern description.

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