[inductor] Fix nan-handling of max and min reductions

[peterbell10]

Stack from ghstack (oldest at bottom):

• [inductor] Reduce duplication of reduction combine functions #99661
• [inductor] Enable conditional use of tl.reduce #100569
• [inductor] Stop using x + tl.zeros(...) in generated triton #100163
• [inductor] Fix argmin/max with duplicate values #99920
• -> [inductor] Fix nan-handling of max and min reductions #99881

This adds helpers that replace tritons minimum, maximum, min and
max with the correct NaN prrpagation. I also removed
ops.int_minimum in favor of ops.minimum because we can just omit
the nan-checks by checking the dtype.

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/99881

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Degradation on most runner types due to networking outage

✅ No Failures

As of commit 22e315e:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ngimel]

does this change generated code/increase overhead? minimum/maximum fns are used pretty often (including in register-sensitive contexts, e.g. when fusing relu to matmul), so we should avoid increasing register pressure.

[peterbell10]

If you mean this exact line then it will get DCE'd. Or do you mean the NaN checks?

[ngimel]

Yeah this line, if it's dce'd it's great

[peterbell10]

Just to substantiate this a bit more, here is the triton IR generated for triton_helper.min:

  tt.func private @"min__fp32S1_4S__1cconstexpr[1]"(%arg0: tensor<1x4xf32>) -> tensor<1xf32> {
    %0 = "tt.reduce"(%arg0) ({
    ^bb0(%arg1: f32, %arg2: f32):
      %1 = tt.call @minimum__fp32_fp32__(%arg1, %arg2) : (f32, f32) -> f32
      tt.reduce.return %1 : f32
    }) {axis = 1 : i32} : (tensor<1x4xf32>) -> tensor<1xf32>
    tt.return %0 : tensor<1xf32>
  }
  tt.func private @minimum__fp32_fp32__(%arg0: f32, %arg1: f32) -> f32 {
    %0 = arith.cmpf olt, %arg0, %arg1 : f32
    %1 = tt.call @is_floating__fp32__(%arg0) : (f32) -> i1
    %2 = scf.if %1 -> (i1) {
      %4 = arith.cmpf une, %arg0, %arg0 : f32
      %5 = arith.ori %0, %4 : i1
      scf.yield %5 : i1
    } else {
      scf.yield %0 : i1
    }
    %3 = arith.select %2, %arg0, %arg1 : f32
    tt.return %3 : f32
  }
  tt.func private @is_floating__fp32__(%arg0: f32) -> i1 {
    %0 = tt.call @"zeros____0cconstexpr[(constexpr[1],)]_1cconstexpr[int1]"() : () -> tensor<1xi1>
    %1 = tt.splat %arg0 : (f32) -> tensor<1xf32>
    %2 = arith.uitofp %0 : tensor<1xi1> to tensor<1xf32>
    %3 = arith.addf %1, %2 : tensor<1xf32>
    %true = arith.constant true
    tt.return %true : i1
  }

Admittedly, it's pretty ugly. However, it is massively simplified after just the inlining pass, which is the very first pass in triton's optimizer.

    %11 = "tt.reduce"(%10) ({
    ^bb0(%arg5: f32, %arg6: f32):
      %16 = arith.cmpf olt, %arg5, %arg6 : f32
      %17 = arith.cmpf une, %arg5, %arg5 : f32
      %18 = arith.ori %16, %17 : i1
      %19 = arith.select %18, %arg5, %arg6 : f32
      tt.reduce.return %19 : f32
    }) {axis = 1 : i32} : (tensor<1x4xf32>) -> tensor<1xf32>

You can see it removed the branch on is_floating and all of the associated code.