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

[Jianhui Li]

================
@@ -604,55 +604,60 @@ 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.
 static SetVector<int64_t>
 adjustUnitDimsWithSliceDims(const SetVector<int64_t> &unitDims,
                             ArrayRef<int64_t> sliceDims) {
-  // Reconstruct parent's non-sliced dimensions
+  // 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()));
----------------
Jianhui-Li wrote:

There is no assumption here. 

Note that the recovered parent space from sliceDims/unitDims is not necessary the actual parent rank.  As long as the parent space rank covers both maximum number of sliceDims and unitDims, the algorithm works. There is no simple way to find the parent rank during the recovery. 

For your example, sliceDims = [1], unitDims = [0], then
The parentSpaceRank becomes 1 + 0 + 1 = 2

Rank 2 covers both sliceDims and unitDims so it is safe to find the adjusted unitDims. 

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