Add xegpu.sg_id_range attribute

Author: nbpatel

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

@@ -315,4 +315,31 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout"> {
   let genVerifyDecl = 1;
 }
 
+def XeGPU_RangeAttr : XeGPUAttr<"Range", "range"> {
+  let summary = [{Specifies a half-open range}];
+  let description = [{
+    `RangeAttr` is an attribute that defines a half-open range [start, end).
+    The range is inclusive of the start value and exclusive of the end value.
+    One usage of this attribute can be for warp specialization.
+    For warp specialization, this attribute can be attached to a scf.if op like
+    ```mlir
+    scf.if %cond {
+      // some operations
+    }{sg_id_range = #xegpu.range<[2, 4]>}
+    ```
+    In this case, the scf.if op will only be executed for subgroup IDs 2 and 3.
+  }];
+
+  let parameters = (ins
+    "IntegerAttr": $start,
+    "IntegerAttr": $end
+  );
+
+  let builders = [
+    AttrBuilder<(ins "int":$start, "int":$end)>
+  ];
+
+  let assemblyFormat = "`<` `[`$start ```,` $end `]``>`";
+}
+
 #endif // MLIR_DIALECT_XEGPU_IR_XEGPUATTRS_TD

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

@@ -175,41 +175,37 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
     }
 
     // Check if there is warp specialization.
-    auto isWarpSpecialized = [](Operation *op, int64_t &startRange,
-                                int64_t &endRange) -> bool {
+    auto isWarpSpecialized = [](Operation *op, int64_t &startOfRange,
+                                int64_t &endOfRange) -> bool {
       Operation *parent = op->getParentOp();
       // Find the outermost scf::IfOp with xegpu.sg_id_range.
       while (parent) {
         if (auto ifOp = dyn_cast<scf::IfOp>(parent)) {
-          if (Attribute attr = ifOp->getAttr("xegpu.sg_id_range")) {
-            if (auto denseAttr = dyn_cast<DenseI32ArrayAttr>(attr)) {
-              auto values = denseAttr.asArrayRef();
-              if (values.size() == 2) {
-                startRange = values[0];
-                endRange = values[1];
-              }
-            }
+          if (auto attr = llvm::dyn_cast_or_null<xegpu::RangeAttr>(
+                  ifOp->getAttr("sg_id_range"))) {
+            startOfRange = attr.getStart().getInt();
+            endOfRange = attr.getEnd().getInt();
             break;
           }
         }
         parent = parent->getParentOp();
       }
-      // Return false if startRange is 0
-      return (startRange > 0 && endRange > startRange);
+      // Return false if startOfRange is 0
+      return (startOfRange > 0 && endOfRange > startOfRange);
     };
 
-    int64_t startRange = -1, endRange = -1;
-    bool warpSpecialized = isWarpSpecialized(op, startRange, endRange);
+    int64_t startOfRange = -1, endOfRange = -1;
+    bool warpSpecialized = isWarpSpecialized(op, startOfRange, endOfRange);
 
     // If warp specialization is detected, adjust the subgroup id accordingly
     Value adjustedSgId = linearSgId;
     if (warpSpecialized) {
-      // Subtract startRange from the original subgroup id to get the adjusted
+      // Subtract startOfRange from the original subgroup id to get the adjusted
       // sg id
-      Value startRangeVal =
-          rewriter.create<arith::ConstantIndexOp>(loc, startRange);
+      Value startOfRangeVal =
+          rewriter.create<arith::ConstantIndexOp>(loc, startOfRange);
       adjustedSgId =
-          rewriter.createOrFold<index::SubOp>(loc, linearSgId, startRangeVal);
+          rewriter.createOrFold<index::SubOp>(loc, linearSgId, startOfRangeVal);
     }
 
     auto deLinearizeSgId =

mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir

@@ -314,7 +314,7 @@ gpu.func @dpas_no_sg_data(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
         %load =  xegpu.load_nd %tdesc
           : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>
           -> vector<256x128xf32>
-    } {xegpu.sg_id_range = array<i32: 0, 1>}
+    } {sg_id_range = #xegpu.range<[0, 1]>}
     %cond3 = arith.cmpi sge, %sg_id, %c1 : index
     %cond4 = arith.cmpi slt, %sg_id, %c2 : index
     %cond5 = arith.andi %cond3, %cond4 : i1
@@ -333,7 +333,7 @@ gpu.func @dpas_no_sg_data(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
           : !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 32], lane_layout = [4, 8], lane_data = [1, 1]>>
           -> vector<128x256xf32>
         %dpas = xegpu.dpas %load_a, %load_b {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], lane_layout = [4, 8], lane_data = [1, 1]>} : vector<256x128xf32>, vector<128x256xf32> -> vector<256x256xf32>
-     }{xegpu.sg_id_range = array<i32: 1, 2>}
+     }{sg_id_range = #xegpu.range<[1, 2]>}
     %cond6 = arith.cmpi sge, %sg_id, %c2 : index
     %cond7 = arith.cmpi slt, %sg_id, %c31 : index
     %cond8 = arith.andi %cond6, %cond7 : i1
@@ -347,7 +347,7 @@ gpu.func @dpas_no_sg_data(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
         : !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>
         -> vector<128x64xf32>
       %exp = math.exp %load {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>} : vector<128x64xf32>
-    }{xegpu.sg_id_range = array<i32: 2, 32>}
+    }{sg_id_range = #xegpu.range<[2, 32]>}
     gpu.return
   }
 
@@ -377,7 +377,7 @@ gpu.func @dpas_no_sg_data(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
           -> vector<128x64xf32>
         %exp = math.exp %ld {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>} : vector<128x64xf32>
     }
-  } {xegpu.sg_id_range = array<i32: 3, 8>}
+  } {sg_id_range = #xegpu.range<[3, 8]>}
   gpu.return
   }
 }