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

[Charitha Saumya]

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

what is 16/2 here?

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