Remove DistributionLevel enum

Author: akroviakov

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

@@ -167,16 +167,6 @@ def XeGPU_FenceScope: I32EnumAttr<"FenceScope",
   let cppNamespace = "::mlir::xegpu";
 }
 
-def XeGPU_WGLevel: I32EnumAttrCase<"WG", 0, "wg">;
-def XeGPU_SGLevel: I32EnumAttrCase<"SG", 1, "sg">;
-def XeGPU_WILevel: I32EnumAttrCase<"WI", 2, "wi">;
-def XeGPU_DistributionLevel: I32EnumAttr<"DistributionLevel",
-      "Specify target level for offsets distribution utility.",
-      [XeGPU_WGLevel, XeGPU_SGLevel, XeGPU_WILevel]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::xegpu";
-}
-
 def XeGPU_FenceScopeAttr:
   EnumAttr<XeGPU_Dialect, XeGPU_FenceScope, "fence_scope"> {
     let summary = [{Describes the scope of fence.
@@ -234,17 +224,17 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
                     "xegpu::DistributeLayoutAttr",
                     "dropInstData">,
     InterfaceMethod<[{Delinearizes a linear ID into its multidimensional
-                      indices based on the effective `level` layout.}],
+                      indices based on the effective layout level.}],
                     "FailureOr<SmallVector<Value>>",
                     "delinearizeId",
-                    (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId, "xegpu::DistributionLevel": $level)>,
+                    (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId)>,
     InterfaceMethod<[{Generates instructions to compute multidimensional offsets for dist units
-                      assigned to a `level` identified by linearId. The shape parameter
-                      represents the higher-level problem size. Each `level` may access
+                      assigned to a level identified by linearId. The shape parameter
+                      represents the higher-level problem size. Each level may access
                       multiple blocks according to round-robin distribution rules.}],
                     "FailureOr<SmallVector<SmallVector<Value>>>",
                     "computeDistributedOffsets",
-                    (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId, "ArrayRef<int64_t>":$shape, "xegpu::DistributionLevel": $level)>,
+                    (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId, "ArrayRef<int64_t>":$shape)>,
     InterfaceMethod</*desc=*/[{Check if this layout can be achieved by applying a transpose
                      to some other layout according to given permutation of (0...n-1).}],
                     /*retTy=*/"bool",
@@ -487,16 +477,16 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
     }
 
     /// Delinearizes a linear ID into its multidimensional indices
-    /// based on the effective `level` layout.
+    /// based on the effective level of the layout.
     FailureOr<SmallVector<Value>>
-    delinearizeId(OpBuilder &builder, Location loc, Value linearId, xegpu::DistributionLevel level);
+    delinearizeId(OpBuilder &builder, Location loc, Value linearId);
 
     /// Generates instructions to compute multidimensional offsets for dist units
-    /// assigned to a `level` identified by linearId. The shape parameter
+    /// assigned to a level identified by linearId. The shape parameter
     /// represents the higher-level problem size. Each `level` may access
     /// multiple blocks according to round-robin distribution rules.
     FailureOr<SmallVector<SmallVector<Value>>>
-    computeDistributedOffsets(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape, xegpu::DistributionLevel level);
+    computeDistributedOffsets(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
 
     /// Check if this is slice of some other layout.
     bool isSliceOf(const xegpu::DistributeLayoutAttr &other) { return false; }
@@ -653,15 +643,15 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
     /// Delinearizes a linear subgroup ID into its multidimensional indices
     /// based on the effective subgroup layout.
     FailureOr<SmallVector<Value>>
-    delinearizeId(OpBuilder &builder, Location loc, Value linearId, xegpu::DistributionLevel level);
+    delinearizeId(OpBuilder &builder, Location loc, Value linearId);
 
     /// Generates instructions to compute multidimensional offsets for blocks
     /// assigned to a subgroup identified by linearId. The shape parameter
     /// represents the workgroup-level problem size. Each subgroup may access
     /// multiple blocks according to round-robin distribution rules.
 
     FailureOr<SmallVector<SmallVector<Value>>>
-    computeDistributedOffsets(OpBuilder &builder, Location loc, Value linearId,ArrayRef<int64_t> shape, xegpu::DistributionLevel level);
+    computeDistributedOffsets(OpBuilder &builder, Location loc, Value linearId,ArrayRef<int64_t> shape);
 
     /// Check if this is slice of some other layout.
     bool isSliceOf(const xegpu::DistributeLayoutAttr &other);

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -273,8 +273,7 @@ LayoutAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
 }
 
 FailureOr<SmallVector<Value>>
-LayoutAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId,
-                          xegpu::DistributionLevel idLevel) {
+LayoutAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId) {
 
   // TODO: handle order attribute
   auto hasDefaultOrder = [&]() {
@@ -285,9 +284,9 @@ LayoutAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId,
   if (!hasDefaultOrder())
     return mlir::emitError(loc, "order attribute is currently not supported.");
   SmallVector<int64_t> layout;
-  if (idLevel == xegpu::DistributionLevel::SG) {
+  if (isForWorkgroup()) {
     layout = getEffectiveSgLayoutAsInt();
-  } else if (idLevel == xegpu::DistributionLevel::WI) {
+  } else if (isForSubgroup()) {
     layout = getEffectiveLaneLayoutAsInt();
   } else {
     return failure();
@@ -304,14 +303,13 @@ LayoutAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId,
 /// LayoutAttr.
 FailureOr<SmallVector<SmallVector<Value>>>
 LayoutAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
-                                      Value linearId, ArrayRef<int64_t> shape,
-                                      xegpu::DistributionLevel targetLevel) {
+                                      Value linearId, ArrayRef<int64_t> shape) {
   SmallVector<int64_t> layout;
   SmallVector<int64_t> subShape;
-  if (targetLevel == DistributionLevel::SG) {
+  if (isForWorkgroup()) {
     layout = getEffectiveSgLayoutAsInt();
     subShape = getEffectiveSgDataAsInt();
-  } else if (targetLevel == DistributionLevel::WI) {
+  } else if (isForSubgroup()) {
     layout = getEffectiveLaneLayoutAsInt();
     subShape = getEffectiveLaneDataAsInt();
   } else {
@@ -325,7 +323,7 @@ LayoutAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
   }
 
   // delinearize Ids
-  auto maybeIds = delinearizeId(builder, loc, linearId, targetLevel);
+  auto maybeIds = delinearizeId(builder, loc, linearId);
   if (failed(maybeIds))
     return failure();
   SmallVector<Value> ids = *maybeIds;
@@ -384,30 +382,28 @@ SliceAttr SliceAttr::flatten() const {
 }
 
 FailureOr<SmallVector<Value>>
-SliceAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId,
-                         xegpu::DistributionLevel level) {
+SliceAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId) {
   SliceAttr attr = flatten();
   auto parent = dyn_cast<LayoutAttr>(attr.getParent());
-  return parent.delinearizeId(builder, loc, linearId, level);
+  return parent.delinearizeId(builder, loc, linearId);
 }
 
 // Implements DistributeLayoutAttr::computeDistributedOffsets to generate
 // instructions for computing multi-dimensional offsets when distributed by
 // LayoutAttr.
 FailureOr<SmallVector<SmallVector<Value>>>
 SliceAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
-                                     Value linearId, ArrayRef<int64_t> shape,
-                                     xegpu::DistributionLevel targetLevel) {
+                                     Value linearId, ArrayRef<int64_t> shape) {
   assert(getRank() == static_cast<int64_t>(shape.size()) && "invalid shape.");
   if (!isForWorkgroup())
     return failure();
 
   SmallVector<int64_t> layout;
   SmallVector<int64_t> subShape;
-  if (targetLevel == DistributionLevel::SG) {
+  if (isForWorkgroup()) {
     layout = getEffectiveSgLayoutAsInt();
     subShape = getEffectiveSgDataAsInt();
-  } else if (targetLevel == DistributionLevel::WI) {
+  } else if (isForSubgroup()) {
     layout = getEffectiveLaneLayoutAsInt();
     subShape = getEffectiveLaneDataAsInt();
   } else {
@@ -422,7 +418,7 @@ SliceAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
   }
 
   // delinearize Ids
-  auto maybeIds = delinearizeId(builder, loc, linearId, targetLevel);
+  auto maybeIds = delinearizeId(builder, loc, linearId);
   if (failed(maybeIds))
     return failure();
 

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

@@ -911,9 +911,8 @@ static SmallVector<Value> computeDistributedOffsetsForMatrixOp(
     PatternRewriter &rewriter, Location loc, xegpu::DistributeLayoutAttr layout,
     Value laneId, ArrayRef<int64_t> payloadShape, ValueRange origOffsets) {
   SmallVector<Value> newOffsets;
-  ;
-  auto maybeDescOffsets = layout.computeDistributedOffsets(
-      rewriter, loc, laneId, payloadShape, xegpu::DistributionLevel::WI);
+  auto maybeDescOffsets =
+      layout.computeDistributedOffsets(rewriter, loc, laneId, payloadShape);
   if (failed(maybeDescOffsets))
     return {};
   assert(maybeDescOffsets.value().size() == 1 &&

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

@@ -114,8 +114,8 @@ genOffsetsList(ConversionPatternRewriter &rewriter, OpType op,
   // Compute the list of subgroup-relative offsets for sub-tensors or sub-memory
   // descriptors to be accessed, based on the layout information.
   ArrayRef<int64_t> wgShape = op.getDataShape();
-  auto maybeDescOffsets = layout.computeDistributedOffsets(
-      rewriter, loc, sgId, wgShape, xegpu::DistributionLevel::SG);
+  auto maybeDescOffsets =
+      layout.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
   if (failed(maybeDescOffsets))
     return failure();
 
@@ -831,8 +831,8 @@ struct WgToSgArithConstantOp : public OpConversionPattern<arith::ConstantOp> {
       // Get subgroup id
       Value sgId =
           gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
-      auto sgOffsets = layout.computeDistributedOffsets(
-          rewriter, loc, sgId, wgShape, xegpu::DistributionLevel::SG);
+      auto sgOffsets =
+          layout.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
       if (failed(sgOffsets))
         return failure();
 
@@ -1053,8 +1053,8 @@ struct WgToSgVectorStepOp : public OpConversionPattern<vector::StepOp> {
 
     Value sgId =
         gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
-    auto sgOffsets = layout.computeDistributedOffsets(
-        rewriter, loc, sgId, wgShape, xegpu::DistributionLevel::SG);
+    auto sgOffsets =
+        layout.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
     if (failed(sgOffsets))
       return failure();
 

mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp

@@ -200,8 +200,8 @@ class TestStepOpPattern : public OpConversionPattern<vector::StepOp> {
 
     Value sgId =
         gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
-    auto maybeOffsets = sliceAttr.computeDistributedOffsets(
-        rewriter, loc, sgId, wgShape, xegpu::DistributionLevel::SG);
+    auto maybeOffsets =
+        sliceAttr.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
     if (failed(maybeOffsets))
       return failure();