Adding support for scheduling the epilogue computation when smem_epilogue parameter is true

[protonu]

Refactoring some code and adding some support for smem_epilogue

[rdspring1]

Here is my transformation helper. It can be used in scheduleMmaResult. Just trying to avoid some merge conflicts.

void HopperMultipleMatmulScheduler::transformLikeMmaOutput(
    TensorView* tv,
    bool is_mma_result) {
  // TODO Add constraints

  auto apply_k_dim_offset = [is_mma_result](int64_t idx) constexpr {
    return (is_mma_result) ? idx - 1 : idx;
  };

  // Original: [..., Mo, No, Mi, Ni]
  tv->split(apply_k_dim_offset(-2), getM(params_->mma_macro));
  tv->split(apply_k_dim_offset(-1), getN(params_->mma_macro));
  // After Split: [..., Mo, No, Mio, Mii, Nio, Nii]
  tv->reorder({{apply_k_dim_offset(-3), apply_k_dim_offset(-2)}});
  // After Reorder: [..., Mo, No, Mio, Nio, Mii, Nii]
  tv->merge(apply_k_dim_offset(-4));
  // After Merge: [..., Mo, No, Mio * Nio, Mii, Nii]
  tv->axis(apply_k_dim_offset(-3))->parallelize(ParallelType::TIDy);
  // After Parallelize: [..., Mo, No, Mio * Nio (TIDy), Mii, Nii]
}

[protonu]

!build