Enable serial grid reduction for split-K in matmul scheduler

[jacobhinkle]

Stacked on #1456.

This simply enables serial reduction in the matmul scheduler when split-K is used.

[jacobhinkle]

See above comment about init dtype.

[jacobhinkle]

I had to change some surrounding code to fix the following error, which was caused by using a Double scalar for the init value of the reduction in the rfactored MmaOp:

  // Allocate global tensor T13
  grid_sync::blockSerializeWait<false, false, true>(&T13[index_utils::maskedOffset<true, true, false>(blockIdx, gridDim)]);
  #pragma unroll
  for(nvfuser_index_t i93 = 0; i93 < 4LL; ++i93) {
    nvfuser_index_t i94;
    i94 = 32LL * i93;
    nvfuser_index_t i95;
    i95 = i35 + (i30 * i93);
    nvfuser_index_t i96;
    i96 = -(16LL * i93);
    #pragma unroll
    for(nvfuser_index_t i97 = 0; i97 < 8LL; ++i97) {
      nvfuser_index_t i98;
      i98 = i94 + (4LL * i97);
      nvfuser_index_t i99;
      i99 = i95 + (8LL * i97);
      bool b100;
      b100 = i41 < (-(8LL * (i97 + nvfuser_zero)));
      bool b101;
      b101 = b43 && b100;
      #pragma unroll
      for(nvfuser_index_t i102 = 0; i102 < 2LL; ++i102) {
        nvfuser_index_t i103;
        i103 = i98 + (2LL * i102);
        nvfuser_index_t i104;
        i104 = i99 + (i36 * i102);
        nvfuser_index_t i105;
        i105 = i102 + nvfuser_zero;
        bool b106;
        b106 = i42 < (i96 - (8LL * i105));
        bool b107;
        b107 = b100 && b106;
        Array<float, 2LL, 2> T7;
        #pragma unroll
        for(nvfuser_index_t i108 = 0; i108 < 2LL; ++i108) {
          T7[i108] = 0.00000000000000000e+00;
        }
        #pragma unroll
        for(nvfuser_index_t i108 = 0; i108 < 2LL; ++i108) {
          // Allocate global tensor T14
          reduction::serialReductionStep(
            T7[i108],
            T12[(i103 + i108)],
            0.00000000000000000e+00,
            T14[(i104 + (i108 + nvfuser_zero))],
            [](float &a, float b) { a = a + b; },
            index_utils::maskedOffset<false, false, true>(blockIdx, gridDim) == 0,
            index_utils::maskedOffset<false, false, true>(blockIdx, gridDim) == index_utils::maskedSize<false, false, true>(gridDim) - 1,
            b107,
            b107);
        }
        if ((b101 && b106)) {
          loadLocalToGlobal<float, /*vec_size=*/2, /*is_volatile=*/false>( &T6[(i99 + (i36 * i105))], &T7[0LL]);
        }
      }
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
  grid_sync::blockSerializeRelease<false, false, true>(&T13[index_utils::maskedOffset<true, true, false>(blockIdx, gridDim)]);
}
}
/*
CUDA NVRTC compile error: __tmp_kernel_none_f0_c0_r0_g0.cu(9791): error: no instance of function template "<unnamed>::reduction::serialReductionStep" matches the argument list
            argument types are: (float, float, double, float, lambda [](float &, float)->void, __nv_bool, __nv_bool, __nv_bool, __nv_bool)
            reduction::serialReductionStep(
            ^

1 error detected in the compilation of "__tmp_kernel_none_f0_c0_r0_g0.cu".
*/

[naoyam]

LGTM. Congrats on enabling split-K!