[Mlir-commits] [mlir] [mlir][xegpu] Raise the reduction vector size to match maxReduceVectorSize (PR #225122)

[Artem Kroviakov]

================
@@ -1269,42 +1269,139 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
   return {laneLayout, laneData, order};
 }
 
-/// Computes the (lane_layout, lane_data) for a multi-reduction's source layout.
-/// Only the innermost two dims are distributed; leading dims are assumed unit.
-/// `subgroupSize` lanes go on one dim; up to `maxReduceVectorSize` elements are
-/// packed into lane_data on the other. To minimize cross-lane reduction, lanes
-/// are spread across a non-reduction dim when possible so the reduction happens
-/// within a lane. inst_data is the element-wise product lane_layout *
-/// lane_data.
+/// Computes a multi_reduction's source layout as (lane_layout, lane_data,
+/// inst_data), where inst_data = lane_layout * lane_data on every dim.
 ///
-/// e.g. with srcShape=[32, 128], subgroupSize=16, maxReduceVectorSize=2:
-///   - Switch: reductionDims=[1] and consumerReductionDims=[] -> lanes move
-///     to the non-reduction dim 0: lane_layout=[16, 1], lane_data=[1, 2].
-///   - Default: reductionDims=[0, 1] (both reduced) -> lanes stay on the
-///     innermost dim: lane_layout=[1, 16], lane_data=[2, 1].
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
-                                  ArrayRef<int64_t> reductionDims,
-                                  int subgroupSize, int64_t maxReduceVectorSize,
-                                  bool verticalLaneLayout = false) {
+/// `maxReduceVectorSize` is the per-lane vector budget.
+///
+/// 1. Take lane_layout and lane_data from the consumer. A consumer sliced over
+///    exactly `reductionDims` covers every source dim, so its parent is reused
+///    as-is. Otherwise it covers only the surviving dims, one entry each, and
+///    each entry is honored on its own dim. Lanes the consumer leaves unclaimed
+///    go to the innermost reduced dim rather than sit idle.
+///
+/// 2. Limit lane_layout by `srcShape`: a dim can spread over no more lanes than
+///    it has elements to give.
+///
+/// 3. Give one dim a per-lane vector, so a lane reduces up to
+///    `maxReduceVectorSize` elements per instruction instead of one.
+///
+///    Which dim: it must be reduced, hold more than one element, and have
+///    lane_layout 1, so a lane owns the whole run it reduces. Among the
+///    innermost two dims, the innermost eligible one is used.
+///
+///    How much: `maxReduceVectorSize / product(lane_data)`, multiplied onto the
+///    dim's existing lane_data rather than replacing it.
+///
+///    Skipped when the dim just outside the innermost two is non-unit: the
+///    source was likely reshaped from 2D -- the scale computation splits one
+///    dim in two -- and no dim of such a source can carry the vector.
+///
+/// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
+/// lane_layout=[1,2,8], lane_data=[1,1,1], inst_data=[1,2,8].
+/// e.g. src=[16,16] reduce [1], consumer lane=[16] ->
+/// lane_layout=[16,1], lane_data=[1,16], inst_data=[16,16].
+/// e.g. src=[16,32] reduce [1], consumer sliced over [1] with parent
+/// lane=[16,1] data=[1,2] -> the parent is reused, then the remaining 16/2 is
+/// multiplied onto dim 1: lane_data=[1,16], inst_data=[16,16].
+static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+                  SmallVector<int64_t>>
+computeReductionLaneLayoutAndData(
+    ArrayRef<int64_t> srcShape, ArrayRef<int64_t> reductionDims,
+    int subgroupSize, int64_t maxReduceVectorSize,
+    xegpu::DistributeLayoutAttr consumerLayout = nullptr) {
   int srcRank = srcShape.size();
   SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
+  auto isReduced = [&](int dim) {
+    return llvm::is_contained(reductionDims, dim);
+  };
 
+  // Rule 1. Only a parent that covers the source rank and spends the whole
+  // subgroup is trusted; one spending less distributes less than the producer
+  // already does.
+  auto consumerSlice = dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+  bool reusedParent = false;
+  if (consumerSlice &&
+      consumerSlice.getDims().asArrayRef().equals(reductionDims)) {
+    auto parent = consumerSlice.getParent();
+    SmallVector<int64_t> parentLaneLayout =
+        parent.getEffectiveLaneLayoutAsInt();
+    SmallVector<int64_t> parentLaneData = parent.getEffectiveLaneDataAsInt();
+    if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
+        static_cast<int>(parentLaneData.size()) == srcRank &&
+        computeProduct(parentLaneLayout) == subgroupSize) {
+      laneLayout = std::move(parentLaneLayout);
+      laneData = std::move(parentLaneData);
+      reusedParent = true;
+    }
+  }
+
+  if (!reusedParent) {
+    SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
+    if (consumerLayout) {
+      consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+      consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+    }
+    int64_t lanesUsed = 1;
+    for (int i = 0, c = 0; i < srcRank; ++i) {
+      if (isReduced(i))
+        continue;
+      if (c < static_cast<int>(consumerLaneLayout.size())) {
+        laneLayout[i] = consumerLaneLayout[c];
+        lanesUsed *= laneLayout[i];
+      }
+      if (c < static_cast<int>(consumerLaneData.size()))
+        laneData[i] = consumerLaneData[c];
+      ++c;
+    }
+
+    // The innermost dim only, and only when reduced: an outer one would take
+    // lane_layout [1,2] to [8,2] against a producer still supplying [1,2], and
+    // a lane_layout-only convert_layout has no lowering.
+    // TODO: spend these lanes on an outer reduced dim once that is lowerable --
+    // they sit idle today. Widening it now, even just to srcRank-2, breaks
+    // reduction_2D_scalar_cross_sg_intra_lane.
+    int64_t spareLanes = std::max<int64_t>(1, subgroupSize / lanesUsed);
+    int innermost = srcRank - 1;
+    if (spareLanes > 1 && isReduced(innermost))
+      laneLayout[innermost] = spareLanes;
+  }
+
+  // Rule 2, applied to a reused parent as much as a derived layout.
+  for (int i = 0; i < srcRank; ++i)
+    laneLayout[i] = std::min(laneLayout[i],
+                             std::max<int64_t>(1, srcShape[i] / laneData[i]));
+
+  // Rule 3. A non-unit dim just outside the innermost two means the source was
----------------
akroviakov wrote:

> was likely reshaped from 2D

The assumption should be stated as such, not as a speculation.

>  where no dim is a sound place for the vector

Any example for this? 

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