use 'op.get*Mutable()' when getting layout attrs

Signed-off-by: dchigarev <[email protected]>

Author: dchigarev

mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp

@@ -636,15 +636,12 @@ 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 = xegpu::getDistributeLayoutAttr(gatherOp.getResult());
     auto layoutIndices =
-        xegpu::getDistributeLayoutAttr(gatherOp->getOpOperand(numOffsets + 1));
-    auto layoutMask =
-        xegpu::getDistributeLayoutAttr(gatherOp->getOpOperand(numOffsets + 2));
+        xegpu::getDistributeLayoutAttr(gatherOp.getIndicesMutable());
+    auto layoutMask = xegpu::getDistributeLayoutAttr(gatherOp.getMaskMutable());
     auto layoutPassThru =
-        xegpu::getDistributeLayoutAttr(gatherOp->getOpOperand(numOffsets + 3));
-
+        xegpu::getDistributeLayoutAttr(gatherOp.getPassThruMutable());
     SmallVector<xegpu::CachePolicyAttr, 3> cacheHints =
         getOpCacheHints(gatherOp);
     auto xeGatherOp = xegpu::LoadGatherOp::create(
@@ -689,13 +686,12 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
         computeOffsets(rewriter, scatterOp, meta.first, meta.second);
     Value flatMemref = memrefToIndexPtr(scatterOp, rewriter);
 
-    auto numOffsets = scatterOp.getOffsets().size();
     auto layoutIndices =
-        xegpu::getDistributeLayoutAttr(scatterOp->getOpOperand(numOffsets + 1));
+        xegpu::getDistributeLayoutAttr(scatterOp.getIndicesMutable());
     auto layoutMask =
-        xegpu::getDistributeLayoutAttr(scatterOp->getOpOperand(numOffsets + 2));
+        xegpu::getDistributeLayoutAttr(scatterOp.getMaskMutable());
     auto layoutVal =
-        xegpu::getDistributeLayoutAttr(scatterOp->getOpOperand(numOffsets + 3));
+        xegpu::getDistributeLayoutAttr(scatterOp.getValueToStoreMutable());
     SmallVector<xegpu::CachePolicyAttr, 3> cacheHints =
         getOpCacheHints(scatterOp);
     auto storeOp = xegpu::StoreScatterOp::create(