[Mlir-commits] [mlir] [MLIR][XEGPU] Add blocking support for scatter ops (PR #144766)
Jianhui Li
llvmlistbot at llvm.org
Wed Jun 18 14:01:04 PDT 2025
================
@@ -295,12 +297,39 @@ void XeGPUBlockingPass::runOnOperation() {
Type elemTy = type.getElementType();
Type newTy;
- if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(type))
- newTy = xegpu::TensorDescType::get(
- ctx, tileShape, elemTy, tdescTy.getEncoding(),
- tdescTy.getLayoutAttr().dropInstData());
- else
+ if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(type)) {
+
+ Attribute encoding = tdescTy.getEncoding();
+ // If the encoding is a ScatterTensorDescAttr, we need to
+ // potentially adjust the chunk size based on the inst_data.
+ if (tdescTy.isScattered()) {
+ auto scatterAttr = tdescTy.getEncodingAsScatterTensorDescAttr();
+ // mlir::dyn_cast<xegpu::ScatterTensorDescAttr>(encoding);
+ int64_t chunkSize = scatterAttr.getChunkSize().getInt();
+
+ if (chunkSize > 1) {
+ int64_t blockedChunkSize = chunkSize;
+ auto instData = tdescTy.getLayoutAttr().getInstData();
+ if (!instData.empty())
+ blockedChunkSize = instData.asArrayRef().back();
+
+ auto chunkSizeAttr = mlir::IntegerAttr::get(
+ mlir::IntegerType::get(ctx, 64), blockedChunkSize);
+
+ // To create a new attribute with a different chunk_size:
+ auto newEncoding = xegpu::ScatterTensorDescAttr::get(
+ ctx, scatterAttr.getMemorySpace(), chunkSizeAttr);
----------------
Jianhui-Li wrote:
fixed
https://github.com/llvm/llvm-project/pull/144766
More information about the Mlir-commits
mailing list