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

[Charitha Saumya]

================
@@ -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()));
----------------
charithaintc wrote:

Code says'
`int64_t parentSpaceRank = maxDim + sliceDims.size() + 1;`

For your example, sliceDims = [1], unitDims = [0]
So in this example maxDim = 1
sliceDim.size() = 1
So parentSpaceRank = 1 + 1 + 1 = 3 

Still now clear how answer is 2 in our explanation. 

Consider another example, 

Assume parent rank is truly 4 in this case (its fine to not know it).

Slice dims = [0, 1, 2], dimsToMap = [1]
then,   parentSpaceRank = 2 + 3 + 1 = 5

This clearly not correct. 

Why not flatten the slice and just get pass the parent rank to this function. that way we know the correct number and its easy to figure out everything. 

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