Question about reduceOp

[yuanfz98]

Hello,

I don't understand why we have such constraint. I thought the conversion from triton.reduce to linalg.reduce is direct and we just need to copy the body to another.

    // Reduction of arbitrary operations isn't supported because using the first
    // element across the reduction dimension requires us to iterate over a
    // subview that skips over each first element.
    if (reductionOps.size() != 1 ||
        !isReductionOpSupported(reductionOps.front())) {
      return op.emitError("Only support lowering reduction with body "
                          "containing one max(i/f) or add(i/f).");
    }

We have seen cases like Issue #11, and from my understanding the ways to solve this problem are :

1. support multiple ops in the body
2. match the pattern and reduce it into single maxOp

I appreciate your help.

[nhat-nguyen]

You're absolutely right that we can clone the body of tt.reduce to linalg.reduce directly, however this means we won't know what kind of reduce we're dealing with and therefore won't be able to generate the initial reduction values.

Without pattern matching and choosing the correct initial values based on the ops' semantics, we have to use the first elements along the reduction axis and perform the reduction on the remaining elements. However, this results in creatings sub-tensors that aren't always multiple of 2s, which are sub-optimal for certain hardware.

We're planning to propose a change to triton IR so that each reduce op has an enum or tag indicating what kind of reduction the op is (only for the common cases). This way we can support all the common cases without having to do error-pattern pattern matching.

I hope this helps!

[nhat-nguyen]

Regarding 2), we do have a canonicalization pattern that attempts to convert sequence of cmp and select into either max or min (see MinMaxConverter). The pattern is intentionally left simplistic, however, in the short term, we can definitely improve that pattern to cover more cases. If you're interested in contributing, that would be a great way to start and we would very much appreciate the help.

[yuanfz98]

I appreciate your reply! Please assign this job to me if you don't mind.

module {
  tt.func public @fold_cmp(%arg0: !tt.ptr<i32>) {
    %cst_0 = arith.constant dense<0> : tensor<4096xi32>
    %21 = "tt.reduce"(%cst_0) <{axis = 0 : i32}> ({
    ^bb0(%arg5: f32, %arg6: f32):
      %36 = arith.cmpf ogt, %arg5, %arg6 : f32
      %37 = arith.cmpf une, %arg5, %arg5 : f32
      %38 = arith.ori %36, %37 : i1
      %39 = arith.select %38, %arg5, %arg6 : f32
      tt.reduce.return %39 : f32
    }) : (tensor<4096xi32>) -> f32
    tt.store %arg0, %21 {cache = 1 : i32, evict = 1 : i32} : i32
    tt.return
  }
}

Haven't noticed MinMaxConverter before, I think it is the starting point. I may create folders of cmpOp (f/i) and fold self != self -> false, then fold arith.ori with %36 and a false to %36 itself, leaving the original MinMaxConverter unchanged.