tree a546f31dd9014f4178677edad25b6fd84dd5de78
parent 79ae407c95715f0be5a78260c3b135090f679637
author Muzammil <55665739+Muzammiluddin-Syed-ECE@users.noreply.github.com> 1748472442 -0400
committer GitHub <noreply@github.com> 1748472442 -0500
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJoN5J6CRC1aQ7uu5UhlAAALK0QAEL+Rw0j8XxdEU//tpZOW02a
 9PunWJeyLpoVw1O7QA0ncFvdEixYfGlBKs3nbX2S1PamunHsAJgoCTw5dk6EZPEL
 jrzhctYpXtr6yT7nBUNUWsYuu4V/5KY+iIbUhFkArzb8tMwcDqnS+QnHyHhvwShg
 ef5IvQrlrC77QTGUCT8asmkBnX9BEWxODp9QQvy8C2CXigdgb6SuMyBXRTJ+Srcy
 0IFBXL3Bp8bTbJn+GmhJ4PJUprSeZT571+5g9Pjtu3CFTQ7NSJz9QvEKbNSXTBZ7
 6wxn2ZuDHA9FLZtJZmfVY9HZbhR6TR4uybWc+3RmNQSHKhTJv82aYColjp8/oHVD
 k/XofJagMBil9cXR534zidbBU74TxT2VgWx/quXgh9SdmVbvvLB84SofJZQ2HTzq
 iex2dK8lU3YFqx4cwKEXst0BwApv9J4LA27Ei/pmWp9+x/2+AzJM0OQFgBshZGMa
 iXzZOOTcbUHOLd6maJJZZIpOWqSy9Fb20EZGvxZ3eD3DXhnl7gPFHDrhvGdPYzUr
 R2PPdZI98haee/WGLdD6Fu5yQq3YqnZc2LkkgbcuUVhfRXLanCzQyznJzjZqb7TT
 yPPdItbAQ0Wg/cYa3/9T9Vw7XPiH1JxxifQqf2OPsAtwnZmuW0hLcALkTNrQwbe3
 iDeyq8Q8rwlk2ZI7J//r
 =qahw
 -----END PGP SIGNATURE-----
 

[mlir][GPU] Fixes subgroup reduce lowering (#141825)

Fixes the final reduction steps which were taken from an implementation
of scan, not reduction, causing lanes earlier in the wave to have
incorrect results due to masking.

Now aligning more closely with triton implementation :
https://github.com/triton-lang/triton/pull/5019

# Hypothetical example
To provide an explanation of the issue with the current implementation,
let's take the simple example of attempting to perform a sum over 64
lanes where the initial values are as follows (first lane has value 1,
and all other lanes have value 0):
```
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
```
When performing a sum reduction over these 64 lanes, in the current
implementation we perform 6 dpp instructions which in sequential order
do the following:
1) sum over clusters of 2 contiguous lanes
2) sum over clusters of 4 contiguous lanes
3) sum over clusters of 8 contiguous lanes
4) sum over an entire row
5) broadcast the result of last lane in each row to the next row and
each lane sums current value with incoming value.
5) broadcast the result of the 32nd lane to last two rows and each lane
sums current value with incoming value.

After step 4) the result for the example above looks like this:

```
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
```

After step 5) the result looks like this:
```
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
```

After step 6) the result looks like this:
```
[4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
```
Note that the correct value here is always 1, yet after the
`dpp.broadcast` ops some lanes have incorrect values. The reason is that
for these incorrect lanes, like lanes 0-15 in step 5, the
`dpp.broadcast` op doesn't provide them incoming values from other
lanes. Instead these lanes are provided either their own values, or 0
(depending on whether `bound_ctrl` is true or false) as values to sum
over, either way these values are stale and these lanes shouldn't be
used in general.

So what this means:
- For a subgroup reduce over 32 lanes (like Step 5), the correct result
is stored in lanes 16 to 31
- For a subgroup reduce over 64 lanes (like Step 6), the correct result
is stored in lanes 32 to 63.

However in the current implementation we do not specifically read the
value from one of the correct lanes when returning a final value. In
some workloads it seems without this specification, the stale value from
the first lane is returned instead.

# Actual failing test
For a specific example of how the current implementation causes issues,
take a look at the IR below which represents an additive reduction over
a dynamic dimension.
```
!matA = tensor<1x?xf16>
!matB = tensor<1xf16>
#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0)>
func.func @only_producer_fusion_multiple_result(%arg0: !matA) -> !matB {
  %cst_1 = arith.constant 0.000000e+00 : f16
  %c2_i64 = arith.constant 2 : i64
  %0 = tensor.empty() : !matB
  %2 = linalg.fill ins(%cst_1 : f16) outs(%0 : !matB) -> !matB
  %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "reduction"]} ins(%arg0 : !matA) outs(%2 : !matB)  {
  ^bb0(%in: f16, %out: f16):
    %7 = arith.addf %in, %out : f16
    linalg.yield %7 : f16
  } -> !matB
  return %4 : !matB
}
```
When provided an input of type `tensor<1x2xf16>` and values `{0, 1}` to
perform the reduction over, the value returned is consistently 4. By the
same analysis done above, this shows that the returned value is coming
from one of these stale lanes and needs to be read instead from one of
the lanes storing the correct result.

Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>