[Blackwell] Fix thrown away load due to wrong wait placement (#9636)

A faulty check line with no -COUNT- was masking bad codegen. I think.

Provisionally, fixes #9635. I am not very familiar with this area but
trying my best! Feel free to close, edit with maintainer access or
whatever.

# New contributor declaration
- [x] I am not making a trivial change, such as fixing a typo in a
comment.

- [x] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- [x] This PR does not need a test because `Changing existing broken
test & codegen`.

- [x] I have not added any `lit` tests.

Author: LunNova

test/Conversion/tritongpu_to_llvm_blackwell.mlir

@@ -1094,7 +1094,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shar
   // CHECK-LABEL: @tensor_memory_ld_red_min_128x256_4_warps
   // CHECK-COUNT-4: tcgen05.ld.red.sync.aligned.32x32b.x64.min.f32
   // CHECK: tcgen05.wait <load>
-  // CHECK-3: llvm.intr.minnum
+  // CHECK-COUNT-3: llvm.intr.minnum
   tt.func public @tensor_memory_ld_red_min_128x256_4_warps() {
     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #blocked_256N_4w>
     %0 = ttng.tmem_alloc %cst_0 {tensor_memory_col_offset = 0 : i32, tensor_memory_row_offset = 0 : i32} : (tensor<128x256xf32, #blocked_256N_4w>) -> !ttg.memdesc<128x256xf32, #tmem_256N, #ttng.tensor_memory, mutable>
@@ -1105,7 +1105,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shar
   // CHECK-LABEL: @tensor_memory_ld_red_max_128x256_4_warps
   // CHECK-COUNT-4: tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32
   // CHECK: tcgen05.wait <load>
-  // CHECK-3: llvm.intr.maxnum
+  // CHECK-COUNT-3: llvm.intr.maxnum
   tt.func public @tensor_memory_ld_red_max_128x256_4_warps() {
     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #blocked_256N_4w>
     %0 = ttng.tmem_alloc %cst_0 {tensor_memory_col_offset = 0 : i32, tensor_memory_row_offset = 0 : i32} : (tensor<128x256xf32, #blocked_256N_4w>) -> !ttg.memdesc<128x256xf32, #tmem_256N, #ttng.tensor_memory, mutable>
@@ -1125,7 +1125,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shar
   // CHECK-LABEL: @tensor_memory_ld_red_min_128x256_4_warps_nan
   // CHECK-COUNT-4: tcgen05.ld.red.sync.aligned.32x32b.x64.min.NaN.f32
   // CHECK: tcgen05.wait <load>
-  // CHECK-3: llvm.intr.minimum
+  // CHECK-COUNT-3: llvm.intr.minimum
   tt.func public @tensor_memory_ld_red_min_128x256_4_warps_nan() {
     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #blocked_256N_4w_nan>
     %0 = ttng.tmem_alloc %cst_0 {tensor_memory_col_offset = 0 : i32, tensor_memory_row_offset = 0 : i32} : (tensor<128x256xf32, #blocked_256N_4w_nan>) -> !ttg.memdesc<128x256xf32, #tmem_256N_nan, #ttng.tensor_memory, mutable>
@@ -1136,7 +1136,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shar
   // CHECK-LABEL: @tensor_memory_ld_red_max_128x256_4_warps_nan
   // CHECK-COUNT-4: tcgen05.ld.red.sync.aligned.32x32b.x64.max.NaN.f32
   // CHECK: tcgen05.wait <load>
-  // CHECK-3: llvm.intr.maximum
+  // CHECK-COUNT-3: llvm.intr.maximum
   tt.func public @tensor_memory_ld_red_max_128x256_4_warps_nan() {
     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #blocked_256N_4w_nan>
     %0 = ttng.tmem_alloc %cst_0 {tensor_memory_col_offset = 0 : i32, tensor_memory_row_offset = 0 : i32} : (tensor<128x256xf32, #blocked_256N_4w_nan>) -> !ttg.memdesc<128x256xf32, #tmem_256N_nan, #ttng.tensor_memory, mutable>

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp

@@ -412,29 +412,6 @@ std::pair<SmallVector<Value>, SmallVector<Value>> lowerTMemLdSt(
     }
   }
 
-  // Combine partial reductions into one value per thread
-  if (redvalVals.size() > 1) {
-    auto isMin = *redOp == TMEMLoadReduceModifier::MIN;
-    auto applyMinMax = [&](Value lhs, Value rhs) {
-      return useNaN ? (isMin ? LLVM::MinimumOp::create(rewriter, loc, lhs, rhs)
-                             : LLVM::MaximumOp::create(rewriter, loc, lhs, rhs))
-                          ->getResult(0)
-                    : (isMin ? LLVM::MinNumOp::create(rewriter, loc, lhs, rhs)
-                             : LLVM::MaxNumOp::create(rewriter, loc, lhs, rhs))
-                          ->getResult(0);
-    };
-    // Use tree reduction: pair up elements at each level
-    while (redvalVals.size() > 1) {
-      SmallVector<Value> reduced;
-      assert(redvalVals.size() % 2 == 0 &&
-             "redvalVals must be a multiple of 2");
-      for (size_t i = 0; i < redvalVals.size(); i += 2) {
-        reduced.push_back(applyMinMax(redvalVals[i], redvalVals[i + 1]));
-      }
-      redvalVals = std::move(reduced);
-    }
-  }
-
   return {resultVals, redvalVals};
 }
 
@@ -519,6 +496,34 @@ static std::pair<SmallVector<Value>, SmallVector<Value>> lowerTMemLdStFromTypes(
                                vals, tmemBase, redOp, useAbs, useNaN);
 }
 
+// Combine partial reductions into one value per thread via tree reduction.
+static void combinePartialReductions(Location loc,
+                                     ConversionPatternRewriter &rewriter,
+                                     SmallVector<Value> &redvalVals,
+                                     TMEMLoadReduceModifier redOp,
+                                     bool useNaN) {
+  if (redvalVals.size() <= 1)
+    return;
+  auto isMin = redOp == TMEMLoadReduceModifier::MIN;
+  auto applyMinMax = [&](Value lhs, Value rhs) {
+    return useNaN ? (isMin ? LLVM::MinimumOp::create(rewriter, loc, lhs, rhs)
+                           : LLVM::MaximumOp::create(rewriter, loc, lhs, rhs))
+                        ->getResult(0)
+                  : (isMin ? LLVM::MinNumOp::create(rewriter, loc, lhs, rhs)
+                           : LLVM::MaxNumOp::create(rewriter, loc, lhs, rhs))
+                        ->getResult(0);
+  };
+  // Use tree reduction: pair up elements at each level
+  while (redvalVals.size() > 1) {
+    SmallVector<Value> reduced;
+    assert(redvalVals.size() % 2 == 0 && "redvalVals must be a multiple of 2");
+    for (size_t i = 0; i < redvalVals.size(); i += 2) {
+      reduced.push_back(applyMinMax(redvalVals[i], redvalVals[i + 1]));
+    }
+    redvalVals = std::move(reduced);
+  }
+}
+
 struct TensorMemoryLoadOpConversion
     : public ConvertOpToLLVMPattern<triton::nvidia_gpu::TMEMLoadOp> {
   using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
@@ -556,6 +561,10 @@ struct TensorMemoryLoadOpConversion
     // Wait insertion could be moved to the TTGIR level if needed.
     NVVM::Tcgen05WaitOp::create(rewriter, loc, NVVM::Tcgen05WaitKind::LOAD);
 
+    // tcgen05.ld.red is async, redval registers aren't valid until the wait
+    if (redOp)
+      combinePartialReductions(loc, rewriter, redvalVals, *redOp, useNaN);
+
     // Handle reduction output if present
     SmallVector<Value> results = {resultStruct};
     if (redOp) {