refactor

Author: chencha3

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

@@ -236,7 +236,7 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
       return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getType().getLayout());
     }
 
-    ArrayRef<int64_t> getDistributeShape() {
+    ArrayRef<int64_t> getDataShape() {
       return getTensorDescShape();
     }
 
@@ -283,7 +283,7 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", []> {
       return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getTensorDescType().getLayout());
     }
 
-    ArrayRef<int64_t> getDistributeShape() {
+    ArrayRef<int64_t> getDataShape() {
       return getTensorDescType().getShape();
     }
 
@@ -381,7 +381,7 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
       return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getTensorDescType().getLayout());
     }
 
-    ArrayRef<int64_t> getDistributeShape() {
+    ArrayRef<int64_t> getDataShape() {
       return getTensorDescType().getShape();
     }
 
@@ -473,7 +473,7 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
       return dyn_cast_if_present<xegpu::DistributeLayoutAttr>(getTensorDescType().getLayout());
     }
 
-    ArrayRef<int64_t> getDistributeShape() {
+    ArrayRef<int64_t> getDataShape() {
       return getTensorDescType().getShape();
     }
 
@@ -1243,7 +1243,7 @@ def XeGPU_LoadMatrixOp: XeGPU_Op<"load_matrix", [MemoryEffects<[MemRead]>,
       return getMixedValues(getConstOffsets(), getOffsets(), getContext());
     }
 
-    ArrayRef<int64_t> getDistributeShape() {
+    ArrayRef<int64_t> getDataShape() {
       return getRes().getType().getShape();
     }
   }];
@@ -1285,7 +1285,7 @@ def XeGPU_StoreMatrixOp: XeGPU_Op<"store_matrix", [MemoryEffects<[MemWrite]>,
       return getMixedValues(getConstOffsets(), getOffsets(), getContext());
     }
 
-    ArrayRef<int64_t> getDistributeShape() {
+    ArrayRef<int64_t> getDataShape() {
       return getData().getType().getShape();
     }
 

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

@@ -36,15 +36,12 @@ namespace {
 
 // Retrieve the RangeAttr if it is specified.
 static xegpu::RangeAttr getRangeSpecAttr(Operation *op) {
-  Operation *parent = op->getParentOp();
+  Operation *parent = op->getParentOfType<scf::IfOp>();
   while (parent) {
-    if (auto ifOp = dyn_cast<scf::IfOp>(parent)) {
-      if (auto attr = llvm::dyn_cast_or_null<xegpu::RangeAttr>(
-              ifOp->getAttr("sg_id_range"))) {
-        return attr;
-      }
-    }
-    parent = parent->getParentOp();
+    if (auto attr = llvm::dyn_cast_if_present<xegpu::RangeAttr>(
+            parent->getAttr("sg_id_range")))
+      return attr;
+    parent = parent->getParentOfType<scf::IfOp>();
   }
   return {};
 }
@@ -115,14 +112,13 @@ genOffsetsList(ConversionPatternRewriter &rewriter, OpType op,
 
   // Compute the list of subgroup-relative offsets for sub-tensors or sub-memory
   // descriptors to be accessed, based on the layout information.
-  ArrayRef<int64_t> wgShape = op.getDistributeShape();
+  ArrayRef<int64_t> wgShape = op.getDataShape();
   auto maybeDescOffsets = layout.getOffsets(rewriter, loc, sgId, wgShape);
   if (failed(maybeDescOffsets))
     return failure();
 
   // Compute the final global offsets for each accessed sub-tensor
   // or sub-memory descriptor.
-  // SmallVector<SmallVector<OpFoldResult>> offsetsList;
   for (const auto &sgOffsets : *maybeDescOffsets) {
     SmallVector<OpFoldResult> newOffsets = xegpu::addWithRightAligned(
         rewriter, loc, getAsOpFoldResult(sgOffsets), origOffsets);
@@ -777,7 +773,7 @@ struct WgToSgLoadMatrixOp : public OpConversionPattern<xegpu::LoadMatrixOp> {
     if (failed(genOffsetsList(rewriter, op, offsetsList)))
       return failure();
 
-    ArrayRef<int64_t> wgShape = op.getDistributeShape();
+    ArrayRef<int64_t> wgShape = op.getDataShape();
     VectorType valueTy = op.getRes().getType();
     Type elemTy = valueTy.getElementType();