fix

Author: charithaintc

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

@@ -327,8 +327,7 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [AllElementTypesMatch<["value", "Tensor
   let hasVerifier = 1;
 }
 
-def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [AllShapesMatch<["value", "TensorDesc"]>,
-                                       AllElementTypesMatch<["value", "TensorDesc"]>]> {
+def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [AllElementTypesMatch<["value", "TensorDesc"]>]> {
   let summary = "stores a n-D block register region back to memory, currently only supports 2D";
 
   let description = [{

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

@@ -170,7 +170,11 @@ getDistributedTensorDescType(xegpu::TensorDescType originalT,
   for (const auto [l, o] : llvm::zip_equal(layout, shape)) {
     if (!divisible(APInt(64, o), APInt(64, l)))
       return failure();
-    distributedShape.push_back(o / l);
+    // Tensor descriptor is distributed only for the scattered case.
+    if (originalT.isScattered())
+      distributedShape.push_back(o / l);
+    else
+      distributedShape.push_back(o);
   }
   xegpu::TensorDescType distributedDescType;
   if (originalT.isScattered()) {