[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:

6 is the intended answer. 

Rather than recovering the exact parent rank, we compute a safe upper bound so that dimsToMap can be adjusted safely. This upper bound is defined as
max(dimsToMap, sliceDims) + 1 + sliceDims.size().

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