[Mlir-commits] [mlir] [mlir][xegpu] cleanup the print format for TensorDesc (PR #149182)
Chao Chen
llvmlistbot at llvm.org
Fri Jul 18 08:20:47 PDT 2025
https://github.com/chencha3 updated https://github.com/llvm/llvm-project/pull/149182
>From 7717fa74ccb9fc7551a3f209294ad77f1cae547b Mon Sep 17 00:00:00 2001
From: Chao Chen <chao.chen at intel.com>
Date: Wed, 16 Jul 2025 18:25:12 +0000
Subject: [PATCH 1/2] change print format
---
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 6 ++++++
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 16 ++++++++++++----
2 files changed, 18 insertions(+), 4 deletions(-)
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/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())
>From 132bcffb8696bb735771b25c44466af5fb9584ed Mon Sep 17 00:00:00 2001
From: Chao Chen <chao.chen at intel.com>
Date: Wed, 16 Jul 2025 19:59:39 +0000
Subject: [PATCH 2/2] clean up and update the test
---
.../mlir/Dialect/XeGPU/IR/XeGPUTypes.td | 53 ++++++++-----------
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 2 +-
.../VectorToXeGPU/load-to-xegpu.mlir | 6 +--
.../VectorToXeGPU/store-to-xegpu.mlir | 6 +--
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 2 +-
.../transfer-write-to-xegpu.mlir | 2 +-
6 files changed, 30 insertions(+), 41 deletions(-)
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/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>
// -----
More information about the Mlir-commits
mailing list