addresses feedbacks

Author: Jianhui-Li

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

@@ -303,8 +303,7 @@ void XeGPUBlockingPass::runOnOperation() {
       // 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);
+        auto scatterAttr = llvm::dyn_cast_if_present<xegpu::ScatterTensorDescAttr>(encoding);        
         int64_t chunkSize = scatterAttr.getChunkSize().getInt();
 
         if (chunkSize > 1) {
@@ -313,12 +312,9 @@ void XeGPUBlockingPass::runOnOperation() {
           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);
+              ctx, scatterAttr.getMemorySpace().getValue(), blockedChunkSize);
 
           encoding = newEncoding;
         }

mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp

@@ -116,12 +116,9 @@ struct TestXeGPUUnrollingPatterns
                 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);
+                    ctx, scatterAttr.getMemorySpace().getValue(), blockedChunkSize);
 
                 encoding = newEncoding;
               }