[mlir][xegpu] cleanup the print format for TensorDesc #149182
+48
−45
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) ChangesPR #145916 changed all of parameters of Full diff: https://github.com/llvm/llvm-project/pull/149182.diff 8 Files Affected:
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 42b5b7a0d4e3f..d022361d1e376 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -64,6 +64,12 @@ def XeGPU_BlockTensorDescAttr: XeGPU_TensorDescAttr<"BlockTensorDesc", "block_td
)>
];
+ let extraClassDeclaration = [{
+ // return true if all fields of the BlockTensorDescAttr are set with
+ // default values.
+ bool hasDefaultsOnly();
+ }];
+
}
def XeGPU_ScatterTensorDescAttr: XeGPU_TensorDescAttr<"ScatterTensorDesc", "scatter_tdesc_attr"> {
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
index 277158ac85409..c3ab8c9a1b73a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
@@ -131,12 +131,12 @@ def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
return llvm::cast<TensorDescType>(cloneWith(getShape(), elementType));
}
- BlockTensorDescAttr getEncodingAsBlockTensorDescAttr() const {
- return llvm::dyn_cast_if_present<BlockTensorDescAttr>(getEncoding());
- }
-
- ScatterTensorDescAttr getEncodingAsScatterTensorDescAttr() const {
- return llvm::dyn_cast_if_present<ScatterTensorDescAttr>(getEncoding());
+ template <typename T,
+ typename = std::enable_if_t<
+ std::is_same_v<T, BlockTensorDescAttr> ||
+ std::is_same_v<T, ScatterTensorDescAttr>>>
+ T getEncodingOfType() const {
+ return llvm::dyn_cast_if_present<T>(getEncoding());
}
LayoutAttr getLayoutAttr() const {
@@ -144,49 +144,38 @@ def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
}
xegpu::MemorySpace getMemorySpace() const {
- auto block_attr = getEncodingAsBlockTensorDescAttr();
- if (block_attr && block_attr.getMemorySpace())
- return block_attr.getMemorySpace().getValue();
+ if (auto attr = getEncodingOfType<BlockTensorDescAttr>())
+ return attr.getMemorySpace().getValue();
- auto scatter_attr = getEncodingAsScatterTensorDescAttr();
- if (scatter_attr && scatter_attr.getMemorySpace())
- return scatter_attr.getMemorySpace().getValue();
+ if (auto attr = getEncodingOfType<ScatterTensorDescAttr>())
+ return attr.getMemorySpace().getValue();
- // return default value
+ llvm_unreachable("invalid encoding");
return MemorySpace::Global;
}
// get the ArrayLength for blocked TensorDesc
int getArrayLength() {
- auto attr = getEncoding();
- auto block_attr = mlir::dyn_cast_if_present<BlockTensorDescAttr>(attr);
- assert((!attr || block_attr) && "invalid on non BlockTensorDescAttr.");
- if (block_attr && block_attr.getArrayLength())
- return block_attr.getArrayLength().getInt();
- // return default value
- return 1;
+ auto attr = getEncodingOfType<BlockTensorDescAttr>();
+ assert(attr && "invalid on non BlockTensorDescAttr.");
+ return attr.getArrayLength().getInt();
}
bool getBoundaryCheck() {
- auto attr = getEncoding();
- auto block_attr = mlir::dyn_cast_if_present<BlockTensorDescAttr>(attr);
- assert((!attr || block_attr) && "invalid on non BlockTensorDescAttr.");
- if (block_attr && block_attr.getBoundaryCheck())
- return block_attr.getBoundaryCheck().getValue();
- // return default value
- return true;
+ auto attr = getEncodingOfType<BlockTensorDescAttr>();
+ assert(attr && "invalid on non BlockTensorDescAttr.");
+ return attr.getBoundaryCheck().getValue();
}
bool isScattered() {
- return bool(getEncodingAsScatterTensorDescAttr());
+ return bool(getEncodingOfType<ScatterTensorDescAttr>());
}
// get the ChunkSize for scattered TensorDesc
int getChunkSizeAsInt() {
- auto attr = getEncoding();
- auto scatter_attr = mlir::dyn_cast_if_present<ScatterTensorDescAttr>(attr);
- assert(scatter_attr && "invalid on non ScatterTensorDescAttr.");
- return scatter_attr.getChunkSizeAsInt();
+ auto attr = getEncodingOfType<ScatterTensorDescAttr>();
+ assert(attr && "invalid on non ScatterTensorDescAttr.");
+ return attr.getChunkSizeAsInt();
}
/// Helper to drop all layout information from the TensorDesc type.
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 642c393cbc2c8..8ab404d52eab4 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -112,6 +112,11 @@ BlockTensorDescAttr BlockTensorDescAttr::get(mlir::MLIRContext *context,
return Base::get(context, scopeAttr, lengthAttr, boundaryAttr);
}
+bool BlockTensorDescAttr::hasDefaultsOnly() {
+ return getMemorySpace().getValue() == xegpu::MemorySpace::Global &&
+ getArrayLength().getInt() == 1 && getBoundaryCheck().getValue();
+}
+
//===----------------------------------------------------------------------===//
// XeGPU_ScatterTensorDescAttr
//===----------------------------------------------------------------------===//
@@ -253,10 +258,11 @@ mlir::Type TensorDescType::parse(::mlir::AsmParser &parser) {
if (parser.parseGreater())
return {};
+ MLIRContext *ctxt = parser.getContext();
return TensorDescType::getChecked(
- [&]() { return parser.emitError(parser.getNameLoc()); },
- parser.getContext(), shape, elementType,
- encoding.value_or(mlir::Attribute()), layout.value_or(mlir::Attribute()));
+ [&]() { return parser.emitError(parser.getNameLoc()); }, ctxt, shape,
+ elementType, encoding.value_or(BlockTensorDescAttr::get(ctxt)),
+ layout.value_or(mlir::Attribute()));
}
void TensorDescType::print(::mlir::AsmPrinter &printer) const {
@@ -273,7 +279,9 @@ void TensorDescType::print(::mlir::AsmPrinter &printer) const {
printer << getElementType();
- if (auto encoding = getEncoding())
+ auto encoding = getEncoding();
+ auto blockAttr = llvm::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
+ if (encoding && (!blockAttr || !blockAttr.hasDefaultsOnly()))
printer << ", " << encoding;
if (auto layout = getLayout())
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 6b85a66a8bd36..b6df1f00c2462 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -54,7 +54,7 @@ mlir::xegpu::getDistributedVectorType(xegpu::TensorDescType tdescTy) {
std::multiplies<int64_t>());
// Case 1: regular loads/stores
- auto scatterAttr = tdescTy.getEncodingAsScatterTensorDescAttr();
+ auto scatterAttr = tdescTy.getEncodingOfType<ScatterTensorDescAttr>();
if (scatterAttr) {
auto chunkSize = scatterAttr.getChunkSize().getInt();
// Verify if the first dimension of the tensor descriptor shape is
diff --git a/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
index 4af7061a4f8a3..fe7bc9f2395de 100644
--- a/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
@@ -30,7 +30,7 @@ func.func @load_2D_vector(%source: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc
// CHECK-SAME: %[[SRC]][%[[OFFSET]], %[[OFFSET]], %[[OFFSET]]]
-// CHECK-SAME: memref<8x16x32xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<8x16x32xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]{{.*}}-> vector<8x16xf32>
// CHECK: return %[[VEC]]
@@ -55,7 +55,7 @@ func.func @load_dynamic_source(%source: memref<?x?x?xf32>,
// CHECK: %[[DIM_0_STRIDE:.+]] = arith.muli %[[DIM_2]], %[[DIM_1]]
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]][%[[OFFSET]], %[[OFFSET]], %[[OFFSET]]]
// CHECK-SAME: [%[[DIM_0]], %[[DIM_1]], %[[DIM_2]]], [%[[DIM_0_STRIDE]], %[[DIM_2]], 1]
-// CHECK-SAME: memref<?x?x?xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<?x?x?xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]{{.*}}-> vector<8x16xf32>
// CHECK: return %[[VEC]]
@@ -73,7 +73,7 @@ func.func @load_out_of_bounds(%source: memref<7x15xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc
// CHECK-SAME: %[[SRC]][%[[OFFSET]], %[[OFFSET]]]
-// CHECK-SAME: memref<7x15xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<7x15xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]{{.*}}-> vector<8x16xf32>
// CHECK: return %[[VEC]]
diff --git a/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
index d68a02b54e967..53b5699e376b3 100644
--- a/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
@@ -32,7 +32,7 @@ func.func @store_2D_vector(%vec: vector<8x16xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc
// CHECK-SAME: %[[SRC]][%[[OFFSET]], %[[OFFSET]], %[[OFFSET]]]
-// CHECK-SAME: memref<8x16x32xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<8x16x32xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]] : vector<8x16xf32>
// -----
@@ -57,7 +57,7 @@ func.func @store_dynamic_source(%vec: vector<8x16xf32>,
// CHECK: %[[DIM_0_STRIDE:.+]] = arith.muli %[[DIM_2]], %[[DIM_1]]
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]][%[[OFFSET]], %[[OFFSET]], %[[OFFSET]]]
// CHECK-SAME: [%[[DIM_0]], %[[DIM_1]], %[[DIM_2]]], [%[[DIM_0_STRIDE]], %[[DIM_2]], 1]
-// CHECK-SAME: memref<?x?x?xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<?x?x?xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]] : vector<8x16xf32>
// -----
@@ -75,7 +75,7 @@ func.func @store_out_of_bounds(%vec: vector<8x16xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc
// CHECK-SAME: %[[SRC]][%[[OFFSET]], %[[OFFSET]]]
-// CHECK-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]] : vector<8x16xf32>
// -----
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index c2f760b29afc4..8fad4af6608fd 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -51,7 +51,7 @@ func.func @load_zero_pad_out_of_bounds(%source: memref<32x64xf32>,
// CHECK-SAME: %[[SRC:.+]]: memref<32x64xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]][%[[OFFSET]], %[[OFFSET]]]
-// CHECK-SAME: memref<32x64xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<32x64xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]{{.*}}-> vector<8x16xf32>
// CHECK: return %[[VEC]]
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index 8de6c2283b37c..58db6d6bb418b 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -80,7 +80,7 @@ func.func @store_out_of_bounds(%vec: vector<8x16xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc
// CHECK-SAME: %[[SRC]][%[[OFFSET]], %[[OFFSET]]]
-// CHECK-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]] : vector<8x16xf32>
// -----
|
|
|
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.
PR #145916 changed all of parameters of
BlockTensorDescAttrtoDefaultValuedParameters. Therefore, when all parameters are default values,BlockTensorDescAttrwill be printed as empty, e.g.,#xegpu.block_tdesc_attr<>, which is not necessary. This PR cleans up the print ofBlockTensorDescAttr. If all of its fields are default values, it will be not printed in TensorDescType.