[Mlir-commits] [mlir] [MLIR][XeGPU] Refactor layout propagation utilities (PR #179016)

[Jianhui Li]

================
@@ -673,18 +816,59 @@ SliceAttr::setUnitDimData(SetVector<int64_t> unitDims) const {
 
 // set the layout for the sepcified unit dims: sg_lane and lane_layout to 1
 DistributeLayoutAttr
-SliceAttr::setUnitDimLayout(SetVector<int64_t> unitDims) const {
+SliceAttr::setUnitDimLayout(SmallVector<int64_t> unitDims) const {
   DistributeLayoutAttr parentLayout = getParent();
 
   ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
 
-  SetVector<int64_t> adjustUnitDims =
+  SmallVector<int64_t> adjustUnitDims =
       mapSlicedDimsToParentSpace(unitDims, sliceDims);
 
   return SliceAttr::get(
       getContext(), parentLayout.setUnitDimLayout(adjustUnitDims), getDims());
 }
 
+// Derive a new layout with sg_data, inst_data and lane_data set to the
+// specified values for the given dimension
+DistributeLayoutAttr SliceAttr::setDimData(int64_t dim, int64_t sgData,
+                                           int64_t instData, int64_t laneData) {
+  ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
+  auto parent = dyn_cast<LayoutAttr>(getParent());
+
+  SmallVector<int64_t> dimSet;
+  dimSet.push_back(dim);
+  SmallVector<int64_t> adjustDims =
+      mapSlicedDimsToParentSpace(dimSet, sliceDims);
----------------
Jianhui-Li wrote:

say dimset = {1}, sliceDims = {1}. 
This function will map the dimSet to parent space which returns adjustDims = {2}.

We can't simply flatten the layout, since it may destroy the nest slice structure.  SetDimData needs to maintain users' nest slice structure and change the root parent's value only. 

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