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

what happens if parent rank is 4. sliceDims = [1], unitDims = [0]. 

The  parentSpaceRank  becomes 1 + 1 + 1 = 3

Are you making some implicit assumption here? I would assert those assumptions here. 

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