add delinearizeSubgroupId interface

Author: chencha3

mlir/include/mlir/Dialect/XeGPU/IR/XeGPU.h

@@ -15,6 +15,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ShapedOpInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"

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

@@ -187,7 +187,11 @@ def LayoutTrait: AttrInterface<"LayoutTrait"> {
                     "getEffectiveSgLayout">,
     InterfaceMethod<"Get the effective sg data",
                     "std::optional<llvm::SmallVector<int>>",
-                    "getEffectiveSgData">
+                    "getEffectiveSgData">,
+    InterfaceMethod<"Delinearize the Subgroup Id",
+                    "FailureOr<SmallVector<Value>>",
+                    "delinearizeSubgroupId",
+                    (ins "Value":$linearId, "Location":$loc, "OpBuilder &": $builder)>
   ];
 }
 
@@ -358,6 +362,10 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [LayoutTrait]> {
         return llvm::to_vector(data.asArrayRef());
       return std::nullopt;
     }
+
+    FailureOr<SmallVector<Value>>
+    delinearizeSubgroupId(Value linearId, Location loc, OpBuilder &builder);
+
   }];
 
   let assemblyFormat = "`<` struct(params) `>`";
@@ -409,6 +417,9 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [LayoutTrait]> {
       return std::nullopt;
     }
 
+    FailureOr<llvm::SmallVector<Value>>
+    delinearizeSubgroupId(Value linearId, Location loc, OpBuilder &builder);
+
     DenseI32ArrayAttr getOrder() const {
       return getParent().getOrder();
     }

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

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/XeGPU/IR/XeGPU.h"
 #include "mlir/Dialect/XeGPU/IR/XeGPUTargetInfo.h"
@@ -211,6 +212,18 @@ LayoutAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
   return success();
 }
 
+FailureOr<SmallVector<Value>>
+LayoutAttr::delinearizeSubgroupId(Value linearId, Location loc,
+                                  OpBuilder &builder) {
+  assert(isWgLayout() && "delinearizeSubgroupId is only available for "
+                         "workgroup-level layout attribute.");
+  auto dims =
+      llvm::map_to_vector(getSgLayout().asArrayRef(), [&](int32_t d) -> Value {
+        return arith::ConstantIndexOp::create(builder, loc, d);
+      });
+  return affine::delinearizeIndex(builder, loc, linearId, dims);
+}
+
 //===----------------------------------------------------------------------===//
 // XeGPU_SliceAttr
 //===----------------------------------------------------------------------===//
@@ -232,6 +245,12 @@ SliceAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
   return success();
 }
 
+FailureOr<SmallVector<Value>>
+SliceAttr::delinearizeSubgroupId(Value linearId, Location loc,
+                                 OpBuilder &builder) {
+  return getParent().delinearizeSubgroupId(linearId, loc, builder);
+}
+
 //===----------------------------------------------------------------------===//
 // XeGPU_TensorDescType
 //===----------------------------------------------------------------------===//

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

@@ -838,7 +838,9 @@ void ConvertLayoutOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
 } // namespace xegpu
 } // namespace mlir
 
+namespace mlir {
 #include <mlir/Dialect/XeGPU/IR/XeGPUAttrInterface.cpp.inc>
+} // namespace mlir
 #include <mlir/Dialect/XeGPU/IR/XeGPUEnums.cpp.inc>
 #define GET_OP_CLASSES
 #include <mlir/Dialect/XeGPU/IR/XeGPU.cpp.inc>

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

@@ -175,7 +175,7 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
     }
 
     auto deLinearizeSgId =
-        affine::delinearizeIndex(rewriter, loc, linearSgId, sgLayoutDim);
+        layout.delinearizeSubgroupId(linearSgId, loc, rewriter);
     if (failed(deLinearizeSgId))
       return failure();
     SmallVector<Value> sgIds = *deLinearizeSgId;