[WIP][DO NOT REVIEW] Schedule epilogue for Hopper - without smem epilogue #3565
Conversation
| d->axis(-1)->parallelize(ParallelType::Vectorize); | ||
| } | ||
| scheduleFusionInputsForEpilogue(); |
There was a problem hiding this comment.
In this path, since we are not using TMA we should respect the params_->supported_vec_size.epilogue parameter to set the vectorization width. See scheduleOutputTensor in this file, which we should either update and use or remove.
In the TMA path (the else branch) we might not be able to use scheduleFusionInputsForEpilogue as is because it is currently propagating the schedule back from the output d, which is scheduled for a TMA store not a vectorized load. We could potentially just do that propagation then check the innermost dims and merge/split/unroll as needed in order to hit the target vectorization for the epilogue.
protonu
force-pushed
the
pbasu_wip_mma_epilogue
branch
2 times, most recently
from
d4e9ae6 to
75c9a86
Compare
protonu
force-pushed
the
pbasu_wip_mma_epilogue
branch
from
75c9a86 to
99aace1
Compare
| @@ -1310,6 +1310,77 @@ TEST_F(MatmulSchedulerTest, EpilogueAlpha) { | |||
| NVF_CHECK(outputs[0].allclose(tref, 0.001, 0.001)); | |||
| } | |||
|
|
|||
| TEST_F(MatmulSchedulerTest, MatmulHopperCastBiasCast) { | |||
There was a problem hiding this comment.
I think this is fine for now but since this test is not Hopper-specific iin the future (once #3440 is trimmed and merged) we could probably remove it and just unguard some other tests like the MatmulOp and LinearOp translation tests which use the same fusion but with a matmul instead of fusedMultiplySum.
No description provided.