address comments

Author: Jianhui-Li

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

@@ -1319,6 +1319,9 @@ def XeGPU_LoadMatrixOp: XeGPU_Op<"load_matrix", [MemoryEffects<[MemRead]>,
     Arguments:
      - `mem_desc`: the memory descriptor identifying the SLM region.
      - `offsets`: the coordinates within the matrix to read from.
+     - `subgroup_block_io`: [optional] An attribute indicating that the operation can be 
+                 lowered to a subgroup block load. When this attribute is present, 
+                 the offsets are subgroup-uniform across all lanes.
      - `layout`: [optional] An attribute for guiding distributions among
                  subgroups and/or work-items. It currently can accept either
                  LayoutAttr or SliceAttr.
@@ -1367,6 +1370,9 @@ def XeGPU_StoreMatrixOp: XeGPU_Op<"store_matrix", [MemoryEffects<[MemWrite]>,
      - `mem_desc`: the memory descriptor specifying the SLM region.
      - `offsets`: the coordinates within the matrix where the data will be written.
      - `data`: the values to be stored in the matrix.
+     - `subgroup_block_io`: [optional] An attribute indicating that the operation can be 
+                 lowered to a subgroup block store. When this attribute is present, 
+                 the offsets are subgroup-uniform across all lanes.     
      - `layout`: [optional] An attribute for guiding distributions among
                  subgroups and/or work-items. It currently can accept either
                  LayoutAttr or SliceAttr.

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

@@ -263,10 +263,10 @@ def XeGPU_MemDesc: XeGPUTypeDef<"MemDesc", "mem_desc", [ShapedTypeInterface], "m
     /// based on the rank and the value of the first stride dimension.
     bool isColMajor() {
       auto dim0 = dyn_cast<IntegerAttr>(getStrideAttr()[0]);
-      return getRank() == 2 && dim0 && dim0.getInt() == 1;
+      return getRank() == 2 && dim0.getInt() == 1;
     }
 
-    // get the Blocking shape for a MemDescType, Which is represented
+    // Get the Blocking shape for a MemDescType, Which is represented
     // as an attribute in MemDescType. By default it is the shape
     // of the mdescTy
     SmallVector<int64_t> getBlockShape() {
@@ -284,16 +284,18 @@ def XeGPU_MemDesc: XeGPUTypeDef<"MemDesc", "mem_desc", [ShapedTypeInterface], "m
     // Get strides as vector of integer. 
     // If it contains block attribute, the strides are blocked strides.
     //
-    // The blocking is applied against the original matrix shape
-    // so that the linear offset is not impacted by the subview.
+    // The blocking is applied to the base matrix shape derived from the
+    // memory descriptor's stride information. If the matrix described by
+    // the memory descriptor is not contiguous, it is assumed that the base
+    // matrix is contiguous and follows the same memory layout.
     //
     // It first computes the original matrix shape using the stride info,
     // then computes the number of blocks in each dimension of original shape,
     // then compute the outer block shape and stride,
     // then combines the inner and outer block shape and stride
-    // e.g. for mem_desc<32x256xf16, @block=[16, 8], @strides=[1, 32]>
+    // e.g. for `mem_desc<32x256xf16, @block=[16, 8], @strides=[1, 32]>`
     // its memory layout tuple is ([2,32,16,8],[128,256,1,16])
-    // for  mem_desc<256x32xf16, @block=[8, 16]> with default @stride[32, 1]
+    // for  `mem_desc<256x32xf16, @block=[8, 16]>` with default @stride[32, 1]
     // its memory layout tuple is ([32,2,8,16],[256,128,16,1])
     SmallVector<int64_t> getStrideShape();
 

mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp

@@ -520,7 +520,7 @@ class CreateMemDescOpPattern final
   matchAndRewrite(xegpu::CreateMemDescOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
 
-    auto resTy = cast<xegpu::MemDescType>(op.getResult().getType());
+    auto resTy = op.getMemDesc();
 
     // Create the result MemRefType with the same shape, element type, and
     // memory space