address comments

Author: Jianhui-Li

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

@@ -242,7 +242,6 @@ def XeGPU_MemDesc: XeGPUTypeDef<"MemDesc", "mem_desc", [ShapedTypeInterface], "m
       if (layout && layout.hasAttr("stride")) {
         return layout.getStrides();
       }
-
       // derive and return default strides
       SmallVector<int64_t> defaultStrides;
       llvm::append_range(defaultStrides, getShape().drop_front());
@@ -251,6 +250,15 @@ def XeGPU_MemDesc: XeGPUTypeDef<"MemDesc", "mem_desc", [ShapedTypeInterface], "m
       return builder.getI64ArrayAttr(defaultStrides);
     }
 
+    ArrayAttr getBlockAttr() {
+      auto layout = getMemLayout();
+      if (layout && layout.hasAttr("block")) {
+        return layout.getBlockAttr();
+      }
+      Builder builder(getContext());
+      return builder.getI64ArrayAttr({});
+    }
+
     /// Heuristic to determine if the MemDesc uses column-major layout,
     /// based on the rank and the value of the first stride dimension.
     bool isColMajor() {
@@ -261,16 +269,14 @@ def XeGPU_MemDesc: XeGPUTypeDef<"MemDesc", "mem_desc", [ShapedTypeInterface], "m
     // 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> getBlockSize() {
+    SmallVector<int64_t> getBlockShape() {
       SmallVector<int64_t> size(getShape());
-      MemLayoutAttr layout = getMemLayout();
-      if (layout && layout.hasAttr("block")) {
-        ArrayAttr attr = layout.getBlockAttr();
+      ArrayAttr blockAttr = getBlockAttr();
+      if (!blockAttr.empty()) {
         size.clear();
-        llvm::for_each(attr, [&](Attribute elem) {
-          if (auto intElem = dyn_cast<IntegerAttr>(elem))
-            size.push_back(intElem.getInt());
-        });
+        for (auto attr : blockAttr.getValue()) {
+          size.push_back(cast<IntegerAttr>(attr).getInt());
+        }
       }
       return size;
     }
@@ -289,7 +295,7 @@ def XeGPU_MemDesc: XeGPUTypeDef<"MemDesc", "mem_desc", [ShapedTypeInterface], "m
     // 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]
     // its memory layout tuple is ([32,2,8,16],[256,128,16,1])
-    SmallVector<int64_t> getStrides();
+    SmallVector<int64_t> getStrideShape();
 
     /// Generates instructions to compute the linearize offset 
     //  if the memory descriptor is blocked, it returns linearize offset based on the blocked layout

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -776,7 +776,7 @@ SmallVector<OpFoldResult> getBlockedOffsets(OpBuilder &builder, Location loc,
 }
 
 // Get strides as vector of integer for MemDesc.
-SmallVector<int64_t> MemDescType::getStrides() {
+SmallVector<int64_t> MemDescType::getStrideShape() {
 
   SmallVector<int64_t> matrixShape(getShape().begin(), getShape().end());
 
@@ -786,7 +786,7 @@ SmallVector<int64_t> MemDescType::getStrides() {
     strides.push_back(cast<IntegerAttr>(attr).getInt());
   }
 
-  SmallVector<int64_t> innerBlkShape = getBlockSize();
+  SmallVector<int64_t> innerBlkShape = getBlockShape();
 
   // get perm from FCD to LCD
   // perm[i] = the dim with i-th smallest stride
@@ -837,8 +837,8 @@ Value MemDescType::getLinearOffsets(OpBuilder &builder, Location loc,
                                     ArrayRef<OpFoldResult> offsets) {
 
   SmallVector<int64_t> matrixShape(getShape().begin(), getShape().end());
-  SmallVector<int64_t> blockShape = getBlockSize();
-  SmallVector<int64_t> strides = getStrides();
+  SmallVector<int64_t> blockShape = getBlockShape();
+  SmallVector<int64_t> strides = getStrideShape();
 
   // blockshape equal to matrixshape means no blocking
   if (llvm::equal(blockShape, matrixShape)) {

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

@@ -201,10 +201,10 @@ IsValidStoreMatrixParams(VectorType dataTy, MemDescType mdescTy,
       return emitError() << "data shape must not exceed mem_desc shape.";
   } else if (dataShape.size() == 1) {
 
-    SmallVector<int64_t> blockSize = mdescTy.getBlockSize();
+    SmallVector<int64_t> blockShape = mdescTy.getBlockShape();
     // if the subgroup_block_io attribute is set,  mdescTy must have block
     // attribute
-    if (subgroup_block_io && !blockSize.size())
+    if (subgroup_block_io && !blockShape.size())
       return emitError() << "mem_desc must have block attribute when "
                             "subgroup_block_io is set.";
     // if the subgroup_block_io attribute is set, the memdesc should be row

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

@@ -942,6 +942,8 @@ struct UnrollLoadMatrixOp : public UnrollPattern<xegpu::LoadMatrixOp> {
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     VectorType valueTy = llvm::dyn_cast<VectorType>(op.getType());
+    assert(valueTy && "the value type must be vector type!");
+
     std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
     if (!targetShape || targetShape->size() != (size_t)valueTy.getRank())
       return failure();
@@ -985,6 +987,7 @@ struct UnrollStoreMatrixOp : public UnrollPattern<xegpu::StoreMatrixOp> {
 
     Location loc = op.getLoc();
     VectorType valueTy = llvm::dyn_cast<VectorType>(op.getData().getType());
+    assert(valueTy && "the value type must be vector type!");
     ArrayRef<int64_t> shape = valueTy.getShape();
     auto layout = dyn_cast<xegpu::LayoutAttr>(op.getLayoutAttr());
 

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

@@ -992,6 +992,7 @@ struct WgToSgLoadMatrixOp : public OpConversionPattern<xegpu::LoadMatrixOp> {
 
     ArrayRef<int64_t> wgShape = op.getDataShape();
     VectorType valueTy = llvm::dyn_cast<VectorType>(op.getRes().getType());
+    assert(valueTy && "the value type must be vector type!");
     Type elemTy = valueTy.getElementType();
 
     xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();

mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir

@@ -198,4 +198,4 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
     gpu.return %1: vector<8xf16>
   }
 
-}
+}