Overall looks good. Adding a standalone test sounds good.

What is the persistent buffer size in the case of inner and outer reductions? The existing routine to calculate the buffer size is, I believe, looks at reduction domains to calculate the minimum required size to do one reduction. Is that the correct size in the case of the combined reduction pattern? It seems that for the inner reduction that seems like an overestimation as not all of the reduction domains are actually reduced for the inner reduction.

I'm also a little concerned if this schedule should be always used for any inner and outer persistent sizes. I was expecting to see some changes to canScheduleRuntime(). Have you checked some extreme cases?

What is the persistent buffer size in the case of inner and outer reductions? The existing routine to calculate the buffer size is, I believe, looks at reduction domains to calculate the minimum required size to do one reduction. Is that the correct size in the case of the combined reduction pattern? It seems that for the inner reduction that seems like an overestimation as not all of the reduction domains are actually reduced for the inner reduction.

The current persistent buffer detection can't detect the buffer for the partial result of outer reduciton, which is actually a large part of the total persistent buffers.
In the current branch:
For input DataType::Half, the persistent buffers are projected to three inputs (dy, x, weight), total size is 3 * sizeof(half) * dim1
For input DataType::Float the persistent buffers are NOT projected, they are xhat and d_xhat, the total size is 2 * sizeof(float) * dim1
I also tried to disable projection for input DataType::Half, it increased from 123 us to 203 us. But if I enforce projection for input DataType::Float, there is a significiant speedup, e.g. for case 2048 x 10240 the time is reduced from 274 us to 207 us, for case 2048 x 1024 the time is reduced from 39 us to 36 us. The reason is because weight is shared across different rows. If we keep it persistent, we don't need to reload it in the iteration over different rows. The projected version needs more registers per thread but it doesn't reduce the occupancy ratio as the all the blocks must be active at the same time for this grid persistent kernel.

So in the revised version, I enforced the projection and the persistent buffers are:

  weight: float T37[((10 * 1) * 4)];
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T47;
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T43;
  for(nvfuser_index_t i378 = 0; i378 < (ceilDiv(2048, ((nvfuser_index_t)gridDim.y))); ++i378) {
    dy: float T33[((10 * 1) * 4)];
    x : float T34[((10 * 1) * 4)];
  }

as a comparation if not projected, 236 registers per thread (248 registers per thread if projected):

  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T38;
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T34;
  for(nvfuser_index_t i323 = 0; i323 < (ceilDiv(2048, ((nvfuser_index_t)gridDim.y))); ++i323) {
    xhat: float T7[((10 * 1) * 4)];
    d_xhat: float T9[((10 * 1) * 4)];
  }

What is the persistent buffer size in the case of inner and outer reductions? The existing routine to calculate the buffer size is, I believe, looks at reduction domains to calculate the minimum required size to do one reduction. Is that the correct size in the case of the combined reduction pattern? It seems that for the inner reduction that seems like an overestimation as not all of the reduction domains are actually reduced for the inner reduction.

The current persistent buffer detection can't detect the buffer for the partial result of outer reduciton, which is actually a large part of the total persistent buffers.

What do you mean it can't detect the buffer? Does it underestimate the buffer size then?

In the current branch: For input DataType::Half, the persistent buffers are projected to three inputs (dy, x, weight), total size is 3 * sizeof(half) * dim1 For input DataType::Float the persistent buffers are NOT projected, they are xhat and d_xhat, the total size is 2 * sizeof(float) * dim1 I also tried to disable projection for input DataType::Half, it increased from 123 us to 203 us. But if I enforce projection for input DataType::Float, there is a significiant speedup, e.g. for case 2048 x 10240 the time is reduced from 274 us to 207 us, for case 2048 x 1024 the time is reduced from 39 us to 36 us. The reason is because weight is shared across different rows. If we keep it persistent, we don't need to reload it in the iteration over different rows. The projected version needs more registers per thread but it doesn't reduce the occupancy ratio as the all the blocks must be active at the same time for this grid persistent kernel.

So, when input is Float and is projected, the weight tensor also becomes persistent and improves the performance? I don't think this is an intended consequence of the buffer projection, so we should understand why and think if it could be made more generic.

I mean when calculating persistent buffer size, the code can't detect the following two buffers. So yes, there is an under estimation.

  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T38;
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T34;

Enforce projection to input can improve performance because weight is shared across different rows. If we keep it persistent, we don't need to reload it in the iteration over different rows.

I mean when calculating persistent buffer size, the code can't detect the following two buffers. So yes, there is an under estimation.

  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T38;
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T34;

Don't we need to account for the buffers to decide if the persistent scheduler should be picked? This is why I also asked:

I'm also a little concerned if this schedule should be always used for any inner and outer persistent sizes. I was expecting to see some changes to canScheduleRuntime(). Have you checked some extreme cases?

Enforce projection to input can improve performance because weight is shared across different rows. If we keep it persistent, we don't need to reload it in the iteration over different rows.

Right, so let me rephrase my question. Is this specific to this inner-outer scheduler? If not, how can we make it more general? Should we always use projected buffers?

I mean when calculating persistent buffer size, the code can't detect the following two buffers. So yes, there is an under estimation.

  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T38;
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T34;

Don't we need to account for the buffers to decide if the persistent scheduler should be picked? This is why I also asked:

I'm also a little concerned if this schedule should be always used for any inner and outer persistent sizes. I was expecting to see some changes to canScheduleRuntime(). Have you checked some extreme cases?

Enforce projection to input can improve performance because weight is shared across different rows. If we keep it persistent, we don't need to reload it in the iteration over different rows.

Right, so let me rephrase my question. Is this specific to this inner-outer scheduler? If not, how can we make it more general? Should we always use projected buffers?

I checked all the benckmarks and there are regressions if hidden size is >= 16K with float. In the revised branch, buffer size of outer reductions is added:
persistent_buffer_size += scheduler_utils::partialReductionBufferSize( outer_reduction_tvs, runtime_info);
To allow the use of this combined approach, available_persistent_buffer_size is increased from half of all the registers to all the 64K registers. This will allow float hidden size <= 14K to pass canScheduleRunTime check. This lead to register spills but the performance is still faster (860 GB/s to 660 GB/s) than the segmented version. To avoid register spill, multi-blocks per row shoud be used. This is probabaly not urgent as popular hidden size is usually <= 10K.

I mean when calculating persistent buffer size, the code can't detect the following two buffers. So yes, there is an under estimation.

  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T38;
  partial_result_of_outer_reduction: Array<float, ((10 * 1) * 4), 4> T34;

Don't we need to account for the buffers to decide if the persistent scheduler should be picked? This is why I also asked:

I'm also a little concerned if this schedule should be always used for any inner and outer persistent sizes. I was expecting to see some changes to canScheduleRuntime(). Have you checked some extreme cases?

Enforce projection to input can improve performance because weight is shared across different rows. If we keep it persistent, we don't need to reload it in the iteration over different rows.

Right, so let me rephrase my question. Is this specific to this inner-outer scheduler? If not, how can we make it more general? Should we always use projected buffers?

I checked all the benckmarks and there are regressions if hidden size is >= 16K with float. In the revised branch, buffer size of outer reductions is added: persistent_buffer_size += scheduler_utils::partialReductionBufferSize( outer_reduction_tvs, runtime_info); To allow the use of this combined approach, available_persistent_buffer_size is increased from half of all the registers to all the 64K registers. This will allow float hidden size <= 14K to pass canScheduleRunTime check. This lead to register spills but the performance is still faster (860 GB/s to 660 GB/s) than the segmented version. To avoid register spill, multi-blocks per row shoud be used. This is probabaly not urgent as popular hidden size is usually <= 10K.

Thanks for checking the performance. The size check makes sense to me. Register usage and its perf impact is difficult to predict, but looks like it's a reasonable heuristic. I'll revisit and review the PR.

Have you also run the benchmarks on other devices such as V100?

I checked all the benckmarks and there are regressions if hidden size is >= 16K with float.

This is for both V100 and A100.

#2400 (comment)

  if (rparams.vectorization_factor_tmp_gmem_read > 1) {
    for (auto tv_pair : cached_outputs) {
      if (tv_pair.second->axis(-1)->getParallelType() !=
          ParallelType::Vectorize) {
        tv_pair.second->axis(-1)->parallelize(ParallelType::Vectorize);
      }
    }
  }

The above code have been removed in the revised version. All the cached_outputs are correctly vectorized if I correctly set vectorization in scheduleReductionCombinedOuter, previously was set to unroll.