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

[Jianhui Li]

================
@@ -685,6 +813,47 @@ SliceAttr::setUnitDimLayout(SetVector<int64_t> unitDims) const {
       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());
+
+  SetVector<int64_t> dimSet;
+  dimSet.insert(dim);
+  SetVector<int64_t> adjustDims = mapSlicedDimsToParentSpace(dimSet, sliceDims);
+
+  return SliceAttr::get(
+      getContext(),
+      parent.setDimData(adjustDims[0], sgData, instData, laneData), getDims());
+}
+
+// Derive a new layout by collapsing groups of dimensions.
+// Each inner array in `dimGroups` specifies a set of dimensions
+// that are collapsed into a single dimension in the derived layout.
+DistributeLayoutAttr
+SliceAttr::collapseDims(SmallVector<SmallVector<int64_t>> dimGroups) const {
+
+  // Map the sliced dims from parent space to collapsed space
+  SmallVector<int64_t> sliceDims = llvm::to_vector(getDims().asArrayRef());
+
+  // go through dimGroups and map each dim from sliced space to parent space
+  SmallVector<SmallVector<int64_t>> adjustedDimGroups;
+  for (const auto &group : dimGroups) {
+    SetVector<int64_t> groupSet(group.begin(), group.end());
----------------
Jianhui-Li wrote:

what specific changes you suggest? I change groupSet to group, but run into compilation error.  

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