Add support for stmatrix in the unit test HopperMatmulTest/HSH_NT_128BSwizzle

[protonu]

This demonstrates the use of stmatrix in a multi-tile hopper matmul.

[jacobhinkle]

Just to be clear, does this mean we must use a macro such that getN(macro) is a multiple of 16?

[protonu]

yes, but I'd phrase as use a tile that can divide the macro.
For the m-dimension this shouldn't be a problem since it's always 64 for Hopper.
For the n-dimension, if it's not a multiple of 16, it has to a multiple of 8, and we'll have to use stmatrix.x2 where st_matrix_n_tile will be 8.

[zasdfgbnm]

Could you also post some perf data here with smem epilogue on? The command I use is posted here: #3279

[protonu]

Observations:

1. On running getBankConflictInfo, I see both cases (mma with and without a stmatrix) have zero bank conflicts. The codes I used was:

  KernelExecutor ke;
  auto launch_constraints = LaunchParams();
  ke.compile(
      &fusion,
      {inputs.first, inputs.second},
      launch_constraints,
      matmul_cparams);

  auto bank_conflict_info =
      getBankConflictInfo(ke.kernel(), launch_constraints);

  if (bank_conflict_info.empty()) {
    debug() << "===== No bank confliction =====" << std::endl;
  } else {
    debug() << "======= Bank confliction =======" << std::endl;
    for (auto info : bank_conflict_info) {
      debug() << "Expr: " << info.first->toString() << std::endl;
      auto conflict = info.second;
      if (conflict.first > 1) {
        debug() << "input conflict: " << conflict.first << " way, ";
      }
      if (conflict.second > 1) {
        debug() << "output conflict: " << conflict.second << " way";
      }
      debug() << std::endl;
    }
    debug() << "================================" << std::endl;
  }
  fusion.printKernel();

2. Performance sees a minor degradation:

Without stmatrix
Time (%) Total Time (ns) Instances Avg (ns) Med (ns) Min (ns) Max (ns) StdDev (ns) Name

---

 36.8           156830          1  156830.0  156830.0    156830    156830          0.0  <unnamed>::nvfuser_none_f0_c0_r0_g0(<unnamed>::Tensor<<unnamed>::__half, (int)3, (int)3>, <unnamed>…
 23.1            98654          1   98654.0   98654.0     98654     98654          0.0  nvjet_hsh_256x128_64x4_1x2_h_bz_coopA_NTT

Where stmatrix was about 63%

With stmatrix this falls to about 59-60%

TODO: Use nsight to look at bank conflict data.

[protonu]

!build