[Mlir-commits] [mlir] [mlir][XeGPU][VectorToXeGPU] Propagate vector layouts to xegpu ops (PR #163071)

[Jianhui Li]

================
@@ -616,16 +638,39 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
         computeOffsets(rewriter, gatherOp, meta.first, meta.second);
     Value flatMemref = memrefToIndexPtr(gatherOp, rewriter);
 
+    auto numOffsets = gatherOp.getOffsets().size();
+    auto layoutRes = mlir::xegpu::getDistributeLayoutAttr(gatherOp.getResult());
+    auto layoutIndices = mlir::xegpu::getDistributeLayoutAttr(
+        gatherOp->getOpOperand(numOffsets + 1));
+    auto layoutMask = mlir::xegpu::getDistributeLayoutAttr(
+        gatherOp->getOpOperand(numOffsets + 2));
+    auto layoutPassThru = mlir::xegpu::getDistributeLayoutAttr(
+        gatherOp->getOpOperand(numOffsets + 3));
+
+    SmallVector<xegpu::CachePolicyAttr, 3> cacheHints{xegpu::CachePolicyAttr{},
+                                                      xegpu::CachePolicyAttr{},
+                                                      xegpu::CachePolicyAttr{}};
+    getOpCacheHints(gatherOp, cacheHints);
     auto xeGatherOp = xegpu::LoadGatherOp::create(
         rewriter, loc, vectorType, flatMemref, localOffsets, gatherOp.getMask(),
         /*chunk_size=*/IntegerAttr{},
-        /*l1_hint=*/xegpu::CachePolicyAttr{},
-        /*l2_hint=*/xegpu::CachePolicyAttr{},
-        /*l3_hint=*/xegpu::CachePolicyAttr{});
+        /*l1_hint=*/cacheHints[0],
+        /*l2_hint=*/cacheHints[1],
+        /*l3_hint=*/cacheHints[2]);
+    mlir::xegpu::setDistributeLayoutAttr(xeGatherOp->getOpResult(0), layoutRes);
+    mlir::xegpu::setDistributeLayoutAttr(xeGatherOp->getOpOperand(1),
+                                         layoutIndices);
+    mlir::xegpu::setDistributeLayoutAttr(xeGatherOp->getOpOperand(2),
+                                         layoutMask);
 
     auto selectOp =
         arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
                                 xeGatherOp.getResult(), gatherOp.getPassThru());
----------------
Jianhui-Li wrote:

I think that we should only set the layout for anchor op, and leaving the propogation to set layout for operands of neighbour ops.
In case the value/operands are from function arguments, the propogation should propogate through functions.  The reality is that the use case we have are mostly inlined, so we didnt' run into any function calls within device kernel. 


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