[Mlir-commits] [mlir] [MLIR][XeGPU] Matrix load/store subgroup distribution (PR #165008)

[Adam Siemieniuk]

================
@@ -906,6 +907,123 @@ struct StoreDistribution final : public gpu::WarpDistributionPattern {
   }
 };
 
+template <class MatrixOp>
+struct MatrixOpDistribution final : public gpu::WarpDistributionPattern {
+  using gpu::WarpDistributionPattern::WarpDistributionPattern;
+  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,
+                                PatternRewriter &rewriter) const override {
+    gpu::YieldOp yield = warpOp.getTerminator();
+    Operation *lastNode = yield->getPrevNode();
+    auto matrixOp = dyn_cast_or_null<MatrixOp>(lastNode);
+    if (!matrixOp)
+      return failure();
+    constexpr bool isLoad{std::is_same_v<MatrixOp, xegpu::LoadMatrixOp>};
+    int operandIdx{-1};
+
+    VectorType sgPayloadTy;
+    VectorType warpResultTy;
+    if constexpr (isLoad) {
----------------
adam-smnk wrote:

nit: It could be easier to read when split into to separate patterns. Judging by the amount of `if constexpr` checks and fact that some variables (like `warpResultTy`) are not used in both paths.

Perhaps, common bits in between could be refactored into helpers?
Not a blocker though. You can leave it as is.

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