use isTransposeOf

Author: charithaintc

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

@@ -191,6 +191,9 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
     InterfaceMethod<"Get the rank of attribute",
                     "int64_t",
                     "getRank">,
+    InterfaceMethod<"Get the order field of the attribute as integer array",
+                    "DenseI32ArrayAttr",
+                    "getOrder">,
     InterfaceMethod<"Get the num of effective subgroups",
                     "int64_t",
                     "getNumSubgroups", (ins), [{
@@ -253,33 +256,40 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
                         seen[ta.value()] = true;
                       }
                       auto checkTranspose = [](ArrayRef<int64_t> dst, ArrayRef<int64_t> src, ArrayRef<int64_t> perm) {
+                        // If both `dst` and `src` are empty, conservatively return true
+                        // here because some layout fields can be empty.
+                        if (dst.empty() && src.empty())
+                          return true;
                         for (const auto &ta : llvm::enumerate(perm)) {
                           if (src[ta.index()] != dst[ta.value()])
                             return false;
                         }
                         return true;
                       };
-                      // check sgLayout
+                      // Check sgLayout
                       if (!checkTranspose($_self.getSgLayoutAsInt(), other.getSgLayoutAsInt(), perm))
                         return false;
-                      // check sgData
+                      // Check sgData
                       if (!checkTranspose($_self.getSgDataAsInt(), other.getSgDataAsInt(), perm))
                         return false;
-                      // check instData
+                      // Check instData
                       if (!checkTranspose($_self.getInstDataAsInt(), other.getInstDataAsInt(), perm))
                         return false;
-                      // check laneLayout
+                      // Check laneLayout
                       if (!checkTranspose($_self.getLaneLayoutAsInt(), other.getLaneLayoutAsInt(), perm))
                         return false;
-                      // check laneData
+                      // Check laneData
                       if (!checkTranspose($_self.getLaneDataAsInt(), other.getLaneDataAsInt(), perm))
                         return false;
+                      // Check order if both sides have order field.
+                      if ($_self.getOrder() && other.getOrder()) {
+                        auto thisOrderAsInt = llvm::to_vector_of<int64_t>($_self.getOrder().asArrayRef());
+                        auto otherOrderAsInt = llvm::to_vector_of<int64_t>(other.getOrder().asArrayRef());
+                        if (!checkTranspose(thisOrderAsInt, otherOrderAsInt, perm))
+                          return false;
+                      }
                       return true;
-                    }]>,
-    InterfaceMethod</*desc=*/[{Check if this layout is a slice of some other layout.}],
-                    /*retTy=*/"bool",
-                    /*methodName=*/"isSliceOf",
-                    /*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other)>
+                    }]>
   ];
 }
 
@@ -481,9 +491,6 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
     FailureOr<SmallVector<SmallVector<Value>>>
     getOffsets(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
 
-    /// Check if this is slice of some other layout.
-    bool isSliceOf(const xegpu::DistributeLayoutAttr &other) { return false; }
-
   }];
 
   let assemblyFormat = "`<` struct(params) `>`";
@@ -645,9 +652,6 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
     FailureOr<SmallVector<SmallVector<Value>>>
     getOffsets(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
 
-    /// Check if this is slice of some other layout.
-    bool isSliceOf(const xegpu::DistributeLayoutAttr &other);
-
   }];
 
   let assemblyFormat = "`<` qualified($parent) `,` `dims` `=` $dims `>`";

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

@@ -410,26 +410,6 @@ SliceAttr::getOffsets(OpBuilder &builder, Location loc, Value linearId,
                                   shape);
 }
 
-bool SliceAttr::isSliceOf(const xegpu::DistributeLayoutAttr &other) {
-  auto flattenedThis = flatten();
-  // If other is a LayoutAttr, just compare directly with parent of
-  // flattenedThis.
-  if (auto otherLayout = dyn_cast<xegpu::LayoutAttr>(other))
-    return flattenedThis.getParent() == otherLayout;
-  // If other is a SliceAttr, flatten it first before comparing.
-  auto flattenedOther = dyn_cast<xegpu::SliceAttr>(other).flatten();
-  // Both must have common parent LayoutAttr.
-  if (flattenedThis.getParent() != flattenedOther.getParent())
-    return false;
-  // otherFlattened's sliced dims must be a subset of flattenedThis's sliced
-  // dims.
-  llvm::SmallDenseSet<int64_t> thisDims(
-      flattenedThis.getDims().asArrayRef().begin(),
-      flattenedThis.getDims().asArrayRef().end());
-  return llvm::all_of(flattenedOther.getDims().asArrayRef(),
-                      [&](int64_t dim) { return thisDims.contains(dim); });
-}
-
 //===----------------------------------------------------------------------===//
 // XeGPU_RangeAttr
 //===----------------------------------------------------------------------===//

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

@@ -1108,24 +1108,15 @@ struct VectorTransposeDistribution final : public gpu::WarpDistributionPattern {
           transposeOp,
           "the source or result vector of the transpose op lacks layout "
           "attribute");
-    SmallVector<int64_t> sourceLaneLayout = sourceLayout.getLaneLayoutAsInt();
-    SmallVector<int64_t> resultLaneLayout = resultLayout.getLaneLayoutAsInt();
-    SmallVector<int64_t> sourceLaneData = sourceLayout.getLaneDataAsInt();
-    SmallVector<int64_t> resultLaneData = resultLayout.getLaneDataAsInt();
-    if (sourceLaneLayout.size() != 2 || resultLaneLayout.size() != 2)
+    if (sourceLayout.getRank() != 2 || resultLayout.getRank() != 2)
       return rewriter.notifyMatchFailure(
           transposeOp, "the source or result vector of the transpose op "
                        "does not have 2D layout");
-    auto is2DTranspose = [](ArrayRef<int64_t> input, ArrayRef<int64_t> output) {
-      return input.size() == 2 && output.size() == 2 && input[0] == output[1] &&
-             input[1] == output[0];
-    };
-
-    if (!is2DTranspose(sourceLaneLayout, resultLaneLayout) ||
-        !is2DTranspose(sourceLaneData, resultLaneData))
+    ArrayRef<int64_t> perm = transposeOp.getPermutation();
+    if (!resultLayout.isTransposeOf(sourceLayout, perm))
       return rewriter.notifyMatchFailure(
           transposeOp,
-          "the source or result vector layouts must be transposes of each "
+          "the source or result vector layouts must be 2D transposes of each "
           "other");
     FailureOr<VectorType> distributedSourceTypeOrFailure =
         getDistVecTypeBasedOnLaneLayout(sourceLayout,
@@ -1141,7 +1132,7 @@ struct VectorTransposeDistribution final : public gpu::WarpDistributionPattern {
     rewriter.setInsertionPointAfter(newWarpOp);
     auto newTransposeOp = vector::TransposeOp::create(
         rewriter, newWarpOp.getLoc(), newWarpOp.getResult(newRetIndices[0]),
-        transposeOp.getPermutation());
+        perm);
     Value distributedVal = newWarpOp.getResult(operandIdx);
     rewriter.replaceAllUsesWith(distributedVal, newTransposeOp.getResult());
     return success();