Templatize mask ops

Author: nbpatel

mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

@@ -1270,15 +1270,15 @@ struct WgToSgVectorTransposeOp
   }
 };
 
-// This pattern distributes the vector.constant_mask ops to work at subgroup
-// level.
-struct WgToSgVectorConstantMaskOp
-    : public OpConversionPattern<vector::ConstantMaskOp> {
-  using OpConversionPattern<vector::ConstantMaskOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(vector::ConstantMaskOp op, OneToNOpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
+// Distribute vector mask ops to work at subgroup level.
+template <typename MaskOpType>
+struct WgToSgVectorMaskOp : public OpConversionPattern<MaskOpType> {
+  using OpConversionPattern<MaskOpType>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      MaskOpType op,
+      typename OpConversionPattern<MaskOpType>::OneToNOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
     xegpu::DistributeLayoutAttr layout =
         xegpu::getDistributeLayoutAttr(op.getResult());
     if (!layout || !layout.isForWorkgroup())
@@ -1288,73 +1288,16 @@ struct WgToSgVectorConstantMaskOp
     VectorType type = op.getResult().getType();
     auto wgShape = type.getShape();
 
-    ArrayRef<int64_t> wgMaskDimSizes = op.getMaskDimSizes();
-
-    // Get subgroup ID.
-    Value sgId =
-        gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
-    auto sgOffsets =
-        layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);
-    if (failed(sgOffsets))
-      return failure();
-
-    SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;
-    VectorType resultType = VectorType::get(sgShape, type.getElementType());
-
-    // In each dimension, each subgroup computes its local mask size as:
-    // min(max(wgMaskSize[d] - offset[d], 0), sgDimSize[d])
-    SmallVector<Value> newCreateMaskOps;
-    for (auto offsetSet : *sgOffsets) {
-      SmallVector<Value> maskOperands;
-
-      for (auto [i, wgMaskSize] : llvm::enumerate(wgMaskDimSizes)) {
-        Value wgMaskSizeVal =
-            arith::ConstantIndexOp::create(rewriter, loc, wgMaskSize);
-        Value dimSizeVal =
-            arith::ConstantIndexOp::create(rewriter, loc, sgShape[i]);
-        Value offset = offsetSet[i];
-        Value adjustedMaskSize =
-            arith::SubIOp::create(rewriter, loc, wgMaskSizeVal, offset);
-        Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
-        Value nonNegative =
-            arith::MaxSIOp::create(rewriter, loc, adjustedMaskSize, zero);
-        Value sgMaskSize =
-            arith::MinSIOp::create(rewriter, loc, nonNegative, dimSizeVal);
-        maskOperands.push_back(sgMaskSize);
+    SmallVector<Value> wgMaskDimSizes;
+    if constexpr (std::is_same_v<MaskOpType, vector::ConstantMaskOp>) {
+      for (int64_t maskSize : op.getMaskDimSizes()) {
+        wgMaskDimSizes.push_back(
+            arith::ConstantIndexOp::create(rewriter, loc, maskSize));
       }
-
-      auto newCreateMaskOp =
-          vector::CreateMaskOp::create(rewriter, loc, resultType, maskOperands);
-      xegpu::setDistributeLayoutAttr(newCreateMaskOp->getResult(0),
-                                     layout.dropSgLayoutAndData());
-      newCreateMaskOps.push_back(newCreateMaskOp.getResult());
+    } else if constexpr (std::is_same_v<MaskOpType, vector::CreateMaskOp>) {
+      wgMaskDimSizes = llvm::to_vector(op.getOperands());
     }
 
-    rewriter.replaceOpWithMultiple(op, {newCreateMaskOps});
-    return success();
-  }
-};
-
-// This pattern distributes the vector.create_mask ops to work at subgroup
-// level.
-struct WgToSgVectorCreateMaskOp
-    : public OpConversionPattern<vector::CreateMaskOp> {
-  using OpConversionPattern<vector::CreateMaskOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(vector::CreateMaskOp op, OneToNOpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    xegpu::DistributeLayoutAttr layout =
-        xegpu::getDistributeLayoutAttr(op.getResult());
-    if (!layout || !layout.isForWorkgroup())
-      return failure();
-
-    Location loc = op.getLoc();
-    VectorType type = op.getResult().getType();
-    auto wgShape = type.getShape();
-
-    auto wgMaskOperands = op.getOperands();
-
     Value sgId =
         gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
     auto sgOffsets =
@@ -1366,17 +1309,17 @@ struct WgToSgVectorCreateMaskOp
     VectorType resultType = VectorType::get(sgShape, type.getElementType());
 
     // In each dimension, each subgroup computes its local mask size as:
-    // min(max(wgMaskSize[d] - offset[d], 0), sgDimSize[d])
+    // min(max(wgMaskDimSize[d] - offset[d], 0), sgDimSize[d])
     SmallVector<Value> newCreateMaskOps;
     for (auto offsetSet : *sgOffsets) {
       SmallVector<Value> maskOperands;
 
-      for (auto [i, wgMaskOperand] : llvm::enumerate(wgMaskOperands)) {
+      for (auto [i, wgMaskDimSize] : llvm::enumerate(wgMaskDimSizes)) {
         Value dimSizeVal =
             arith::ConstantIndexOp::create(rewriter, loc, sgShape[i]);
         Value offset = offsetSet[i];
         Value adjustedMaskSize =
-            arith::SubIOp::create(rewriter, loc, wgMaskOperand, offset);
+            arith::SubIOp::create(rewriter, loc, wgMaskDimSize, offset);
         Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
         Value nonNegative =
             arith::MaxSIOp::create(rewriter, loc, adjustedMaskSize, zero);
@@ -1397,6 +1340,8 @@ struct WgToSgVectorCreateMaskOp
   }
 };
 
+using WgToSgVectorConstantMaskOp = WgToSgVectorMaskOp<vector::ConstantMaskOp>;
+using WgToSgVectorCreateMaskOp = WgToSgVectorMaskOp<vector::CreateMaskOp>;
 } // namespace
 
 namespace mlir {