[mlir][XeGPU] Update utils for LayoutAttr and SliceAttr support #154819
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-gpu Author: Chao Chen (chencha3) Changes
To avoid confusion caused by the previous naming conventions, this PR renames and refactors relevant XeGPU utilities to clearly indicate their focus on This PR doesn't change the definition of Patch is 31.84 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/154819.diff 10 Files Affected:
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index b4d696444cc44..77e3c257f234e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -185,6 +185,9 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
InterfaceMethod<"Check the availability of workgroup level layouts",
"bool",
"isForWorkgroup">,
+ InterfaceMethod<"Check the availability of subgroup level layouts",
+ "bool",
+ "isForSubgroup">,
InterfaceMethod<"Get the rank of attribute",
"int64_t",
"getRank">,
@@ -202,9 +205,21 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
InterfaceMethod<"Get the SgData field of the attribute as integer array",
"std::optional<SmallVector<int64_t>>",
"getSgDataAsInt">,
+ InterfaceMethod<"Get the InstData field of the attribute as integer array",
+ "std::optional<SmallVector<int64_t>>",
+ "getInstDataAsInt">,
+ InterfaceMethod<"Get the LaneLayout field of the attribute as integer array",
+ "std::optional<SmallVector<int64_t>>",
+ "getLaneLayoutAsInt">,
+ InterfaceMethod<"Get the LaneData field of the attribute as integer array",
+ "std::optional<SmallVector<int64_t>>",
+ "getLaneDataAsInt">,
InterfaceMethod<"Derive a new layout by dropping sgLayout and sgData",
"xegpu::DistributeLayoutAttr",
"dropSgLayoutAndData">,
+ InterfaceMethod<"Derive a new layout by dropping InstData",
+ "xegpu::DistributeLayoutAttr",
+ "dropInstData">,
InterfaceMethod<[{Delinearizes a linear subgroup ID into its multidimensional
indices based on the effective subgroup layout.}],
"FailureOr<SmallVector<Value>>",
@@ -388,6 +403,24 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
return std::nullopt;
}
+ std::optional<SmallVector<int64_t>> getInstDataAsInt() const {
+ if (DenseI32ArrayAttr inst = getInstData())
+ return llvm::to_vector_of<int64_t>(inst.asArrayRef());
+ return std::nullopt;
+ }
+
+ std::optional<SmallVector<int64_t>> getLaneLayoutAsInt() const {
+ if (DenseI32ArrayAttr layout = getLaneLayout())
+ return llvm::to_vector_of<int64_t>(layout.asArrayRef());
+ return std::nullopt;
+ }
+
+ std::optional<SmallVector<int64_t>> getLaneDataAsInt() const {
+ if (DenseI32ArrayAttr data = getLaneData())
+ return llvm::to_vector_of<int64_t>(data.asArrayRef());
+ return std::nullopt;
+ }
+
/// Delinearizes a linear subgroup ID into its multidimensional indices
/// based on the effective subgroup layout.
FailureOr<SmallVector<Value>>
@@ -488,6 +521,42 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
return std::nullopt;
}
+ /// Returns the InstData of the attribute, computed by applying
+ /// the slice dimensions to the underlying LayoutAttr.
+ std::optional<SmallVector<int64_t>> getInstDataAsInt() const {
+ SliceAttr attr = flatten();
+ auto parent = dyn_cast<LayoutAttr>(attr.getParent());
+ if (auto inst = parent.getInstDataAsInt()) {
+ ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
+ return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(*inst), dims);
+ }
+ return std::nullopt;
+ }
+
+ /// Returns the LaneLayout of the attribute, computed by applying
+ /// the slice dimensions to the underlying LayoutAttr.
+ std::optional<SmallVector<int64_t>> getLaneLayoutAsInt() const {
+ SliceAttr attr = flatten();
+ auto parent = dyn_cast<LayoutAttr>(attr.getParent());
+ if (auto layout = parent.getLaneLayoutAsInt()) {
+ ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
+ return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(*layout), dims);
+ }
+ return std::nullopt;
+ }
+
+ /// Returns the LaneData of the attribute, computed by applying
+ /// the slice dimensions to the underlying LayoutAttr.
+ std::optional<SmallVector<int64_t>> getLaneDataAsInt() const {
+ SliceAttr attr = flatten();
+ auto parent = dyn_cast<LayoutAttr>(attr.getParent());
+ if (auto data = parent.getLaneDataAsInt()) {
+ ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
+ return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(*data), dims);
+ }
+ return std::nullopt;
+ }
+
SliceAttr dropSgLayoutAndData() {
SliceAttr attr = flatten();
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
index 76d58e5ea2424..c173b93face98 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
@@ -40,7 +40,7 @@ def XeGPU_Dialect : Dialect {
let extraClassDeclaration = [{
/// Checks if the given shape can be evenly distributed based on the layout
/// and data factors provided by the LayoutAttr.
- static bool isEvenlyDistributable(llvm::ArrayRef<int64_t> shape, xegpu::LayoutAttr attr);
+ static bool isEvenlyDistributable(llvm::ArrayRef<int64_t> shape, xegpu::DistributeLayoutAttr attr);
/// drops/slices the shape in the specified dims, and return the rest. e.g.,
/// for shape = [32, 64, 8], dims = [0, 2], it will return [64]
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index ab471a1f33ef9..2f6671c5e37cc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1162,8 +1162,8 @@ def XeGPU_ConvertLayoutOp: XeGPU_Op<"convert_layout", [Pure, AllTypesMatch<["sou
the IR is lowered to WI level because that is the end result of all distributions.
}];
let arguments = (ins XeGPU_VectorType: $source,
- XeGPU_LayoutAttr: $input_layout,
- XeGPU_LayoutAttr: $target_layout);
+ DistributeLayoutAttr: $input_layout,
+ DistributeLayoutAttr: $target_layout);
let results = (outs XeGPU_VectorType: $result);
let assemblyFormat = [{
$source prop-dict attr-dict `:` type($source)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index b2b2d3ab85231..82fd70571c022 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -21,6 +21,7 @@ class ValueRange;
class TypeConverter;
namespace xegpu {
+class DistributeLayoutAttr;
class LayoutAttr;
class TensorDescType;
} // namespace xegpu
@@ -60,22 +61,33 @@ FailureOr<VectorType> getDistributedVectorType(xegpu::TensorDescType tdescTy);
FailureOr<VectorType> getDistributedVectorType(VectorType originalType,
LayoutAttr layout);
-/// Return the attribute name for the OpOperand to attach LayoutAttr
+/// Return the attribute name for the OpOperand to attach DistributeLayoutAttr
std::string getLayoutName(const OpOperand &operand);
-/// Return the attribute name for the OpResult to attach LayoutAttr
+/// Return the attribute name for the OpResult to attach DistributeLayoutAttr
std::string getLayoutName(const OpResult result);
-/// Retrieves the LayoutAttr associated with a given Value. For TensorDescType
-/// values, the LayoutAttr is extracted from the TensorDescType itself. For
-/// other values, it is obtained from the attributes of the defining operation.
-/// Returns nullptr if no LayoutAttr is found.
-LayoutAttr getLayoutAttr(const Value value);
+/// Retrieves the DistributeLayoutAttr associated with a given Value. For
+/// TensorDescType values, the DistributeLayoutAttr is extracted from the
+/// TensorDescType itself. For other values, it is obtained from the attributes
+/// of the defining operation. Returns nullptr if no DistributeLayoutAttr is
+/// found.
+DistributeLayoutAttr getDistributeLayoutAttr(const Value value);
-/// Retrieves the LayoutAttr associated with a given OpOperand. It will
-/// first check the operand_layout_{id} of the owner operation. If not found,
-/// it will check the operand itself and its defining op.
-LayoutAttr getLayoutAttr(const OpOperand &opr);
+template <typename AttrTy>
+AttrTy getDistributeLayoutAttrOfType(const Value value) {
+ return dyn_cast_if_present<AttrTy>(getDistributeLayoutAttr(value));
+}
+
+/// Retrieves the DistributeLayoutAttr associated with a given OpOperand. It
+/// will first check the operand_layout_{id} of the owner operation. If not
+/// found, it will check the operand itself and its defining op.
+DistributeLayoutAttr getDistributeLayoutAttr(const OpOperand &opr);
+
+template <typename AttrTy>
+AttrTy getDistributeLayoutAttrOfType(const OpOperand &opr) {
+ return dyn_cast_if_present<AttrTy>(getDistributeLayoutAttr(opr));
+}
/// Removes the LayoutAttr for a given OpOperand or OpResult if it exists.
template <typename T,
@@ -83,23 +95,24 @@ template <typename T,
std::is_same_v<T, OpResult>>>
void removeLayoutAttr(const T &operandOrResult);
-/// Removes the LayoutAttr for each OpOperand and OpResult of the given
-/// operation if they exist. If the operation contains regions, it is also
+/// Removes the DistributeLayoutAttr for each OpOperand and OpResult of the
+/// given operation if they exist. If the operation contains regions, it is also
/// applied recursively to the contained operations
void removeLayoutAttrs(Operation *op);
-/// Sets the LayoutAttr for a given OpOperand or OpResult by attaching
+/// Sets the DistributeLayoutAttr for a given OpOperand or OpResult by attaching
/// it to the owner's dictionary attributes
template <typename T,
typename = std::enable_if_t<std::is_same_v<T, OpOperand> ||
std::is_same_v<T, OpResult>>>
-void setLayoutAttr(const T &operandOrResult, const LayoutAttr layout);
-
-/// Set the LayoutAttr for each OpOperand and OpResult of the given operation.
-/// If the operation contains regions, it is also applied recursively to the
-/// contained operations
-void setLayoutAttrs(Operation *op,
- function_ref<LayoutAttr(Value)> getLayoutImpl);
+void setDistributeLayoutAttr(const T &operandOrResult,
+ const DistributeLayoutAttr layout);
+
+/// Set the DistributeLayoutAttr for each OpOperand and OpResult of the given
+/// operation. If the operation contains regions, it is also applied recursively
+/// to the contained operations
+void setDistributeLayoutAttrs(
+ Operation *op, function_ref<DistributeLayoutAttr(Value)> getLayoutImpl);
/// Extract a set of small vectors from a value with a given shape using
/// vector.extract_stride_slice
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index a2d708be0e937..6de6049facfc6 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -91,7 +91,7 @@ genOffsetsComputingInsts(OpBuilder &builder, Location loc,
// Checks if the given shape can be evenly distributed based on the layout
// and data factors provided by the LayoutAttr.
bool XeGPUDialect::isEvenlyDistributable(llvm::ArrayRef<int64_t> shape,
- xegpu::LayoutAttr attr) {
+ xegpu::DistributeLayoutAttr attr) {
assert(attr && "Layout attribute is missing.");
// Checks whether the given shape can be evenly distributed using the
@@ -104,52 +104,51 @@ bool XeGPUDialect::isEvenlyDistributable(llvm::ArrayRef<int64_t> shape,
// smaller than `layout[i] * data[i]`, allowing multiple compute units to
// share the data.
auto tryDistribute = [&](llvm::ArrayRef<int64_t> shape,
- DenseI32ArrayAttr layout, DenseI32ArrayAttr data,
+ std::optional<SmallVector<int64_t>> layout,
+ std::optional<SmallVector<int64_t>> data,
bool rr = true) -> optional<SmallVector<int64_t>> {
llvm::SmallVector<int64_t> newShape(shape);
if (layout) {
- auto vec = llvm::to_vector_of<int64_t>(layout.asArrayRef());
- if (vec.size() != shape.size())
+ if ((*layout).size() != shape.size())
return std::nullopt;
- auto ratio = computeShapeRatio(shape, vec);
+ auto ratio = computeShapeRatio(shape, *layout);
if (!ratio.has_value())
return std::nullopt;
newShape = ratio.value();
}
if (data) {
- auto vec = llvm::to_vector_of<int64_t>(data.asArrayRef());
- if (vec.size() != shape.size())
+ if ((*data).size() != shape.size())
return std::nullopt;
- auto ratio = computeShapeRatio(newShape, vec);
+ auto ratio = computeShapeRatio(newShape, *data);
if (!ratio.has_value() && rr)
- ratio = computeShapeRatio(vec, newShape);
+ ratio = computeShapeRatio(*data, newShape);
if (!ratio.has_value())
return std::nullopt;
// if data is not null, we always return it for next phase.
- newShape = vec;
+ newShape = *data;
}
return newShape;
};
// check the sgLayout and sgData
auto maybeSgShape =
- tryDistribute(shape, attr.getSgLayout(), attr.getSgData());
+ tryDistribute(shape, attr.getSgLayoutAsInt(), attr.getSgDataAsInt());
if (!maybeSgShape)
return false;
auto sgShape = maybeSgShape.value();
// check InstData, it neither have layout nor need round-robin
auto maybeInstShape =
- tryDistribute(sgShape, nullptr, attr.getInstData(), false);
+ tryDistribute(sgShape, std::nullopt, attr.getInstDataAsInt(), false);
if (!maybeInstShape)
return false;
auto instShape = maybeInstShape.value();
// check LaneLayout and LaneData
- auto maybeLaneShape =
- tryDistribute(instShape, attr.getLaneLayout(), attr.getLaneData(), false);
+ auto maybeLaneShape = tryDistribute(instShape, attr.getLaneLayoutAsInt(),
+ attr.getLaneDataAsInt(), false);
return maybeLaneShape.has_value();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index b3144e4c1e55d..80e9d4d25b06c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -84,9 +84,9 @@ struct ConvertLayoutOpPattern
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
PatternRewriter &rewriter) const override {
- xegpu::LayoutAttr input_layout = op.getInputLayoutAttr();
- xegpu::LayoutAttr target_layout = op.getTargetLayoutAttr();
- if (!input_layout.getInstData() || !target_layout.getInstData())
+ xegpu::DistributeLayoutAttr input_layout = op.getInputLayoutAttr();
+ xegpu::DistributeLayoutAttr target_layout = op.getTargetLayoutAttr();
+ if (!input_layout.getInstDataAsInt() || !target_layout.getInstDataAsInt())
return rewriter.notifyMatchFailure(op, "Not a target ConvertLayoutOp.");
input_layout = input_layout.dropInstData();
@@ -140,10 +140,11 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
else
value = (Value)operandOrResult;
- xegpu::LayoutAttr layout = xegpu::getLayoutAttr(operandOrResult);
+ xegpu::DistributeLayoutAttr layout =
+ xegpu::getDistributeLayoutAttr(operandOrResult);
if (layout && layout.isForSubgroup()) {
- if (auto inst_data = layout.getInstData())
- return llvm::to_vector_of<int64_t>(inst_data.asArrayRef());
+ if (auto inst_data = layout.getInstDataAsInt())
+ return inst_data.value();
if (auto type = dyn_cast<ShapedType>(value.getType()))
return llvm::to_vector(type.getShape());
@@ -204,12 +205,14 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
// skip the op if any of its operands or results has workgroup level layouts
bool hasWgLayoutOperands =
llvm::any_of(op->getOpOperands(), [](OpOperand &opr) {
- xegpu::LayoutAttr layout = xegpu::getLayoutAttr(opr);
+ xegpu::DistributeLayoutAttr layout =
+ xegpu::getDistributeLayoutAttr(opr);
return layout && layout.isForWorkgroup();
});
bool hasWgLayoutResults =
llvm::any_of(op->getOpResults(), [](OpResult result) {
- xegpu::LayoutAttr layout = xegpu::getLayoutAttr(result);
+ xegpu::DistributeLayoutAttr layout =
+ xegpu::getDistributeLayoutAttr(result);
return layout && layout.isForWorkgroup();
});
if (hasWgLayoutOperands || hasWgLayoutResults) {
@@ -220,8 +223,8 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
auto isUnrollable = [](Value value, ArrayRef<int64_t> tileShape) {
Type valTy = value.getType();
if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(valTy)) {
- xegpu::LayoutAttr layout = tdescTy.getLayoutAttr();
- return layout && layout.getInstData();
+ xegpu::DistributeLayoutAttr layout = tdescTy.getLayoutAttr();
+ return layout && layout.getInstDataAsInt();
}
auto shapedType = dyn_cast<ShapedType>(valTy);
return shapedType && !llvm::equal(tileShape, shapedType.getShape());
@@ -247,7 +250,8 @@ void XeGPUBlockingPass::runOnOperation() {
// Preserve the LayoutAttr for each operand to the owner's DictionaryAttr.
// This ensures that the LayoutAttr remains accessible even if the defining
// operation is replaced.
- xegpu::setLayoutAttrs(op, [](Value v) { return xegpu::getLayoutAttr(v); });
+ xegpu::setDistributeLayoutAttrs(
+ op, [](Value v) { return xegpu::getDistributeLayoutAttr(v); });
auto getTileShapeAndCount = [](llvm::ArrayRef<int64_t> shape,
xegpu::LayoutAttr layout) {
@@ -377,7 +381,7 @@ void XeGPUBlockingPass::runOnOperation() {
if (auto layout = op->getAttrOfType<xegpu::LayoutAttr>(name)) {
op->removeAttr(name);
if (!isa<LoopLikeOpInterface>(op))
- xegpu::setLayoutAttr(result, layout.dropInstData());
+ xegpu::setDistributeLayoutAttr(result, layout.dropInstData());
}
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index bef88042fc663..5cb47b2accd68 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -718,7 +718,7 @@ static LogicalResult updateOp(mlir::OpBuilder &builder, mlir::Operation *op,
}
// If the result is a vector type, add a temporary layout attribute to the
// op.
- xegpu::setLayoutAttr(result, layout);
+ xegpu::setDistributeLayoutAttr(result, layout);
}
return success();
}
@@ -800,7 +800,7 @@ updateControlFlowOps(mlir::OpBuilder &builder,
// If the type is a vector type and this region argument is an OpResult,
// set the layout attribute on the OpResult.
if (auto result = dyn_cast<OpResult>(successorInput))
- xegpu::setLayoutAttr(result, successorOperandLayout);
+ xegpu::setDistributeLayoutAttr(result, successorOperandLayout);
}
}
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 2088c3c7fc5ec..e48e2180197ec 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -841,14 +841,15 @@ void XeGPUSubgroupDistributePass::runOnOperation() {
if (!isa<VectorType>(operand.get().getType()))
continue;
- xegpu::LayoutAttr layout = xegpu::getLayoutAttr(operand);
+ auto layout =
+ xegpu::getDistributeLayoutAttrOfType<xegpu...
[truncated]
|
There was a problem hiding this comment.
Pull Request Overview
This PR refactors XeGPU utilities to distinguish between two types of layout attributes introduced with recent changes: DistributeLayoutAttr (including LayoutAttr and SliceAttr) for layout distribution, and MemLayoutAttr for memory layout description. The refactor renames functions and variables to clearly indicate their focus on DistributeLayoutAttr while adding support for SliceAttr.
Key changes:
- Renamed utility functions from
getLayoutAttrtogetDistributeLayoutAttrto clarify purpose - Updated function signatures to use
DistributeLayoutAttrinstead ofLayoutAttrwhere appropriate - Added template functions for type-specific layout attribute retrieval
- Modified transform passes to use the updated utility functions
Reviewed Changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | Renamed core utility functions and updated function signatures to use DistributeLayoutAttr |
| mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | Updated workgroup-to-subgroup distribution pass to use new utility function names |
| mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp | Updated subgroup distribution pass with new utility functions and added template usage |
| mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | Updated layout propagation pass to use renamed utility functions |
| mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp | Updated blocking pass to use DistributeLayoutAttr and new utility functions |
| mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | Updated dialect functions to use DistributeLayoutAttr and new method calls |
| mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | Updated header with new function signatures and template functions |
| mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | Updated ConvertLayoutOp to use DistributeLayoutAttr |
| mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td | Updated dialect declaration to use DistributeLayoutAttr |
| mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | Added new interface methods for accessing layout data as integers |
| @@ -882,7 +883,8 @@ void XeGPUSubgroupDistributePass::runOnOperation() { | |||
| if (vecRank == 0) | |||
| return AffineMap::get(val.getContext()); | |||
| // Get the layout of the vector type. | |||
| xegpu::LayoutAttr layout = xegpu::getLayoutAttr(val); | |||
| // TODO: support more layout types | |||
There was a problem hiding this comment.
The TODO comment should be more specific about which layout types need to be supported and provide context about the current limitation.
Suggested change
| // TODO: support more layout types | |
| // TODO: Currently, only simple lane layouts where each distributed dimension | |
| // corresponds directly to a dimension in the lane_layout attribute are supported. | |
| // More complex layout types (e.g., blocked, permuted, or tiled layouts) are not | |
| // handled. To support these, this function would need to interpret the layout | |
| // attribute accordingly and compute the distributed dimensions based on the | |
| // specific layout semantics. |
Copilot uses AI. Check for mistakes.
| @@ -611,8 +613,9 @@ struct WgToSgConvertLayoutOp | |||
| LogicalResult | |||
| matchAndRewrite(xegpu::ConvertLayoutOp op, OneToNOpAdaptor adaptor, | |||
| ConversionPatternRewriter &rewriter) const override { | |||
| xegpu::LayoutAttr input = op.getInputLayout(); | |||
| xegpu::LayoutAttr target = op.getTargetLayout(); | |||
| // TODO: currently, we only support LayoutAttr | |||
There was a problem hiding this comment.
The TODO comment should explain why only LayoutAttr is currently supported and what needs to be done to support other DistributeLayoutAttr types like SliceAttr.
Suggested change
| // TODO: currently, we only support LayoutAttr | |
| // TODO: currently, we only support LayoutAttr because the logic below assumes | |
| // that both input and target layouts are of type LayoutAttr, which provides | |
| // methods like getSgLayout(), getSgData(), and getOrder(). Other types of | |
| // DistributeLayoutAttr, such as SliceAttr, are not handled here and would | |
| // require additional logic to extract the relevant layout information and | |
| // perform the appropriate conversion. To support other DistributeLayoutAttr | |
| // types like SliceAttr, extend this pattern to handle those types, ensuring | |
| // that the necessary properties are accessed and the conversion semantics | |
| // are correctly implemented for each attribute type. |
Copilot uses AI. Check for mistakes.
charithaintc
approved these changes
Aug 26, 2025
| @@ -185,6 +185,9 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> { | |||
| InterfaceMethod<"Check the availability of workgroup level layouts", | |||
| "bool", | |||
| "isForWorkgroup">, | |||
| InterfaceMethod<"Check the availability of subgroup level layouts", | |||
| "bool", | |||
| "isForSubgroup">, | |||
There was a problem hiding this comment.
nit: hasSgLayouts express the intenstion better.
There was a problem hiding this comment.
I rolled back the naming after thinking it twice. I don't want to mess up it with sgLayout field.
silee2
reviewed
Aug 26, 2025
| @@ -388,6 +403,24 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> { | |||
| return std::nullopt; | |||
| } | |||
|
|
|||
| std::optional<SmallVector<int64_t>> getInstDataAsInt() const { | |||
There was a problem hiding this comment.
Just curious.
Do we have cases of valid empty SmallVector?
If not, checking for std::nullopt could be replaced with a size check?
There was a problem hiding this comment.
I feel it is a good suggestion.
There was a problem hiding this comment.
updated them to return SmallVector<int64_t> with size 0.
silee2
approved these changes
Aug 26, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
SliceAttrwas introduced in #150146, but the XeGPU utilities don't support it yet. Additionally, the arrival ofMemLayoutAttr(#153273) introduces two distinct attribute categories in XeGPU:DistributeLayoutAttr(which includesLayoutAttrandSliceAttr) for layout distribution,MemLayoutAttrfor memory layout description.To avoid confusion caused by the previous naming conventions, this PR renames and refactors relevant XeGPU utilities to clearly indicate their focus on
DistributeLayoutAttr, while also adding support forSliceAttr. It also updated the WgToSg and Blocking passes based on this change.This PR doesn't change the definition of
TensorDescTypeand implementation of subgroup-distribution pass, which still takeLayoutAttronly.