refactor

Author: chencha3

mlir/include/mlir/Dialect/XeGPU/IR/XeGPU.h

@@ -22,18 +22,18 @@
 namespace mlir {
 namespace xegpu {
 class TensorDescType;
+class LayoutAttr;
 } // namespace xegpu
 } // namespace mlir
 
+#include <mlir/Dialect/XeGPU/IR/XeGPUDialect.h.inc>
 #include <mlir/Dialect/XeGPU/IR/XeGPUAttrInterface.h.inc>
 #include <mlir/Dialect/XeGPU/IR/XeGPUEnums.h.inc>
+
 #define GET_ATTRDEF_CLASSES
 #include <mlir/Dialect/XeGPU/IR/XeGPUAttrs.h.inc>
 #define GET_TYPEDEF_CLASSES
 #include <mlir/Dialect/XeGPU/IR/XeGPUTypes.h.inc>
-
-#include <mlir/Dialect/XeGPU/IR/XeGPUDialect.h.inc>
-
 #define GET_OP_CLASSES
 #include <mlir/Dialect/XeGPU/IR/XeGPU.h.inc>
 

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

@@ -396,24 +396,15 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [LayoutTrait]> {
     std::optional<llvm::SmallVector<int32_t>> getEffectiveSgLayout() const {
       if (DenseI32ArrayAttr layout = getParent().getSgLayout()) {
         llvm::ArrayRef<int64_t> dims = getDims().asArrayRef();
-        llvm::SmallVector<int32_t> result;
-        for (auto [i, v]: llvm::enumerate(layout.asArrayRef())) {
-          if (!llvm::is_contained(dims, i))
-            result.push_back(v);
-        }
-        return result;
+        return XeGPUDialect::dropDims(layout.asArrayRef(), dims);
       }
       return std::nullopt;
     }
+
     std::optional<llvm::SmallVector<int32_t>> getEffectiveSgData() const {
       if (DenseI32ArrayAttr data = getParent().getSgData()) {
         llvm::ArrayRef<int64_t> dims = getDims().asArrayRef();
-        llvm::SmallVector<int32_t> result;
-        for (auto [i, v]: llvm::enumerate(data.asArrayRef())) {
-          if (!llvm::is_contained(dims, i))
-            result.push_back(v);
-        }
-        return result;
+        return XeGPUDialect::dropDims(data.asArrayRef(), dims);
       }
       return std::nullopt;
     }

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

@@ -41,6 +41,18 @@ def XeGPU_Dialect : Dialect {
       /// Checks if the given shape can be evenly distributed based on the layout
       /// and data factors provided by the LayoutAttr.
       static bool isEvenlyDistributable(llvm::ArrayRef<int64_t> shape, xegpu::LayoutAttr attr);
+
+      /// drops the data in the specified dimension, and return the rest. e.g.,
+      /// for data = [32, 64, 8], dropPositions = [0, 2], it will return [64]
+      template<typename T, typename U>
+      static llvm::SmallVector<T> dropDims(llvm::ArrayRef<T> data, llvm::ArrayRef<U> dropPositions) {
+        llvm::SmallVector<T> result;
+        for (auto [i, v]: llvm::enumerate(data)) {
+          if (!llvm::is_contained(dropPositions, i))
+            result.push_back(v);
+        }
+        return result;
+      }
     }];
 }