[Mlir-commits] [mlir] [MLIR][XeGPU] setUnitDim bug fix and add documentation (PR #173521)

[llvmlistbot at llvm.org]

github-actions[bot] wrote:

<!--LLVM CODE FORMAT COMMENT: {clang-format}-->


:warning: C/C++ code formatter, clang-format found issues in your code. :warning:

<details>
<summary>
You can test this locally with the following command:
</summary>

``````````bash
git-clang-format --diff origin/main HEAD --extensions cpp -- mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp --diff_from_common_commit
``````````

:warning:
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing `origin/main` to the base branch/commit you want to compare against.
:warning:

</details>

<details>
<summary>
View the diff from clang-format here.
</summary>

``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 13bf47592..d96322e6f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -400,7 +400,8 @@ bool LayoutAttr::isEqualTo(const xegpu::DistributeLayoutAttr &other) {
 }
 
 // set the layout for unit dims: sg_data, inst_data and lane_data to 1
-DistributeLayoutAttr LayoutAttr::setUnitDimData(SetVector<int64_t> unitDims) const{
+DistributeLayoutAttr
+LayoutAttr::setUnitDimData(SetVector<int64_t> unitDims) const {
   auto sgDataOpt = getSgData();
   auto instDataOpt = getInstData();
   auto laneDataOpt = getLaneData();
@@ -441,7 +442,8 @@ DistributeLayoutAttr LayoutAttr::setUnitDimData(SetVector<int64_t> unitDims) con
 }
 
 // set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
-DistributeLayoutAttr LayoutAttr::setUnitDimLayout(SetVector<int64_t> unitDims) const{
+DistributeLayoutAttr
+LayoutAttr::setUnitDimLayout(SetVector<int64_t> unitDims) const {
   auto sgLayoutOpt = getSgLayout();
   auto laneLayoutOpt = getLaneLayout();
 
@@ -587,7 +589,7 @@ bool SliceAttr::isSliceOf(const xegpu::DistributeLayoutAttr &other) {
   // dims.
   llvm::SmallDenseSet<int64_t> thisDims(
       flattenedThis.getDims().asArrayRef().begin(),
-      flattenedThis.getDims().asArrayRef().end()); 
+      flattenedThis.getDims().asArrayRef().end());
   return llvm::all_of(flattenedOther.getDims().asArrayRef(),
                       [&](int64_t dim) { return thisDims.contains(dim); });
 }
@@ -605,19 +607,22 @@ bool SliceAttr::isEqualTo(const xegpu::DistributeLayoutAttr &other) {
 
 // Helper function to adjust unit dimensions from sliced space to parent space
 // say we have a parent shape of rank 4, and slice dims [1,3], so the sliced
-// shape is of rank 2, if we want to set unit dim [0] in sliced space, it maps to
-// dim [0] in parent space; if we want to set unit dim [1] in sliced space, it maps to
-// dim [2] in parent space.
+// shape is of rank 2, if we want to set unit dim [0] in sliced space, it maps
+// to dim [0] in parent space; if we want to set unit dim [1] in sliced space,
+// it maps to dim [2] in parent space.
 static SetVector<int64_t>
 adjustUnitDimsWithSliceDims(const SetVector<int64_t> &unitDims,
                             ArrayRef<int64_t> sliceDims) {
   // get max number from sliceDims and unitDims to determine parent space rank
   int64_t maxDim = -1;
-  maxDim = std::max(maxDim, *std::max_element(sliceDims.begin(), sliceDims.end()));
-  maxDim = std::max(maxDim, *std::max_element(unitDims.begin(), unitDims.end()));
+  maxDim =
+      std::max(maxDim, *std::max_element(sliceDims.begin(), sliceDims.end()));
+  maxDim =
+      std::max(maxDim, *std::max_element(unitDims.begin(), unitDims.end()));
   int64_t parentSpaceRank = maxDim + sliceDims.size() + 1;
 
-  // get remaining dims in parent space after applying slicing with parent's slice Dims
+  // get remaining dims in parent space after applying slicing with parent's
+  // slice Dims
   llvm::SmallDenseSet<int64_t> slicedDimsSet(sliceDims.begin(),
                                              sliceDims.end());
   SmallVector<int64_t> remainingDims;
@@ -629,35 +634,39 @@ adjustUnitDimsWithSliceDims(const SetVector<int64_t> &unitDims,
   // Map unit dims from sliced space to parent space
   SetVector<int64_t> adjustUnitDims;
   for (auto dim : unitDims) {
-      int64_t mappedDim = remainingDims[dim];
-      adjustUnitDims.insert(mappedDim);
+    int64_t mappedDim = remainingDims[dim];
+    adjustUnitDims.insert(mappedDim);
   }
 
   return adjustUnitDims;
 }
 
 // set the layout for unit dims: sg_data, inst_data and lane_data to 1
-DistributeLayoutAttr SliceAttr::setUnitDimData(SetVector<int64_t> unitDims) const{
+DistributeLayoutAttr
+SliceAttr::setUnitDimData(SetVector<int64_t> unitDims) const {
   DistributeLayoutAttr parentLayout = getParent();
 
   ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
 
-  SetVector<int64_t> adjustUnitDims = adjustUnitDimsWithSliceDims(unitDims, sliceDims);
+  SetVector<int64_t> adjustUnitDims =
+      adjustUnitDimsWithSliceDims(unitDims, sliceDims);
 
-  return SliceAttr::get(getContext(), parentLayout.setUnitDimData(adjustUnitDims),
-                        getDims());
+  return SliceAttr::get(getContext(),
+                        parentLayout.setUnitDimData(adjustUnitDims), getDims());
 }
 
 // set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
-DistributeLayoutAttr SliceAttr::setUnitDimLayout(SetVector<int64_t> unitDims) const{
+DistributeLayoutAttr
+SliceAttr::setUnitDimLayout(SetVector<int64_t> unitDims) const {
   DistributeLayoutAttr parentLayout = getParent();
 
   ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
 
-  SetVector<int64_t> adjustUnitDims = adjustUnitDimsWithSliceDims(unitDims, sliceDims);
+  SetVector<int64_t> adjustUnitDims =
+      adjustUnitDimsWithSliceDims(unitDims, sliceDims);
 
-  return SliceAttr::get(getContext(), parentLayout.setUnitDimLayout(adjustUnitDims),
-                        getDims());
+  return SliceAttr::get(
+      getContext(), parentLayout.setUnitDimLayout(adjustUnitDims), getDims());
 }
 
 //===----------------------------------------------------------------------===//

``````````

</details>


https://github.com/llvm/llvm-project/pull/173521