Grouped grid welford

[naoyam]

Enables grouping of grid welford ops across iterations. Same functionality as the iteration grouping for GridReduction. This ins intended to improve the outer-norm grid persistence in batchnorm-like fusions.

To support the grouping, added GroupedWelfordOp and kir::GroupedGridWelford. In the case of grouping non-welford reductions, we support both grouping of multiple separate reduction expressions and grouping of the same expression across iterations. In the case of welford, I only enabled iteration grouping as I don't know if it's useful to be able to group multiple welford expressions. Specifically, groupReductions() is the function to group multiple ReductionOp exprs, which isn't extended to support welford. Iteration grouping is done with IterationDomain::parallelize(ParallelType::Group).

There are a lot of mechanical changes as I added new classes GroupedWelfordOp and kir::GroupedGridWelford. Along the addition of the two classes, I modified some of the WelfordOp and kir::GridWelford code, which I think should be cleaner than before.

As the runtime implementation of the grouped grid welford, ParallelReduce::welfordGroup in fused_reduction.cu is added. It is mostly just a copy of ParalleReduce::reduceGroup but supports welford reductions. Thought about extending reduceGroup for welford as well, but that would lead to even more complicated template designs, which I'm not really comfortable.

[naoyam]

Tried to clean the code up a little with a class to represent the three parameters of welford.

[naoyam]

Ordering of out, in and init seems to vary, which can be confusing and bug prone. Tried to make it always out, in and then init for consistency.

[naoyam]

All the changes in this file are mostly cleanup, except those newly added for GroupedWelfordOp and GroupedGridWelford. The cleanup is meant to make the printing more consistently done, whereas previously the KIR classes had been formatted rather differently.

[naoyam]

Code snippet from FusionCrossIterationGroupedGridAllreduceWelford1:

Without grouping:

  fused_reduction::ParallelReduce<1, 0, 3, 1, 0, 3, true, true> T2_reduction;
  #pragma unroll
  for(nvfuser_index_t i46 = 0; i46 < 2; ++i46) {
    int64_t i177;
    i177 = (((((nvfuser_index_t)blockIdx.x) * 32) + ((nvfuser_index_t)threadIdx.x)) * 2) + (i46 + nvfuser_zero);
    T2_reduction.reduce(
      RefTuple<float, float, int64_t>(T2[i46], T3[i46], T4[i46]),
      ConstRefTuple<float, float, int64_t>(T1[i46], 0, 1),
      VolatilePtrTuple<float, float, int64_t>(&T7[0], &T8[0], &T9[0]),
      &T10[0],
      PtrTuple<float, float, int64_t>(reinterpret_cast<float*>(shared_mem_avg), reinterpret_cast<float*>(shared_mem_var), reinterpret_cast<int64_t*>(shared_mem_n)),
      ((i65 < T0.size[0]) && (i177 < T0.size[1])),
      (i177 < T0.size[1]),
      LocalTuple<float, float, int64_t>(0.00000000000000000e+00, 0.00000000000000000e+00, 0),
      welfordCombine<float, int64_t>);
  }

With grouping:

  fused_reduction::ParallelReduce<1, 0, 3, 1, 0, 3, true, true> T2_reduction;
  T2_reduction.welfordGroup<2, float, int64_t>(RefTuple<float, float>(T2[0], T2[1]),
    RefTuple<float, float>(T3[0], T3[1]),
    RefTuple<int64_t, int64_t>(T4[0], T4[1]),
    ConstRefTuple<float, float>(T1[0], T1[1]),
    ConstRefTuple<float, float>(0, 0),
    ConstRefTuple<int64_t, int64_t>(1, 1),
    LocalTuple<float, float>(0.00000000000000000e+00, 0.00000000000000000e+00),
    LocalTuple<float, float>(0.00000000000000000e+00, 0.00000000000000000e+00),
    LocalTuple<int64_t, int64_t>(0, 0),
    VolatilePtrTuple<float, float>(&T7[0], &T7[((((nvfuser_index_t)gridDim.x) * ((nvfuser_index_t)gridDim.y)) * 32) * 1]),
    VolatilePtrTuple<float, float>(&T8[0], &T8[((((nvfuser_index_t)gridDim.x) * ((nvfuser_index_t)gridDim.y)) * 32) * 1]),
    VolatilePtrTuple<int64_t, int64_t>(&T9[0], &T9[((((nvfuser_index_t)gridDim.x) * ((nvfuser_index_t)gridDim.y)) * 32) * 1]),
    &T10[0],
    PtrTuple<float, float, int64_t>(reinterpret_cast<float*>(shared_mem_avg), reinterpret_cast<float*>(shared_mem_var), reinterpret_cast<int64_t*>(shared_mem_n)),
    LocalTuple<bool, bool>(((i65 < T0.size[0]) && (((((((nvfuser_index_t)blockIdx.x) * 32) + ((nvfuser_index_t)threadIdx.x)) * 2) + (0 + nvfuser_zero)) < T0.size[1])), ((i65 < T0.size[0]) && (((((((nvfuser_index_t)blockIdx.x) * 32) + ((nvfuser_index_t)threadIdx.x)) * 2) + (1 + nvfuser_zero)) < T0.size[1]))),
    LocalTuple<bool, bool>((((((((nvfuser_index_t)blockIdx.x) * 32) + ((nvfuser_index_t)threadIdx.x)) * 2) + (0 + nvfuser_zero)) < T0.size[1]), (((((((nvfuser_index_t)blockIdx.x) * 32) + ((nvfuser_index_t)threadIdx.x)) * 2) + (1 + nvfuser_zero)) < T0.size[1])));

[csarofeen]

I don't know if it's useful to be able to group multiple welford expressions

I can't think of any cases off the top of my head.

[csarofeen]

I didn't see any obvious errors. LGTM

[csarofeen]

Are there any GPUs we frequently run tests on that this wouldn't fit on?

[naoyam]

Forgot to add a check. Thanks for the comment.

[csarofeen]

Didn't check this schedule carefully, I assume once outer grid persistence is in we'll just send it to the segmenter/schedulers.