Schedule epilogue (for Hopper Matmul) by propagation backward from output - smem epilogue not supported. #3580

protonu · 2024-12-12T14:49:52Z

This adds support for scheduling epilogue for the hopper matmul scheduler.
We don't support smem epilogue as yet.

We also don't honor the vectorization_factor as yet for the store to output. That'll be covered in a separate PR.

tests/cpp/test_matmul_scheduler.cpp

jacobhinkle

LGTM. Only comments are around code organization.

csrc/scheduler/hopper_multi_matmul.cpp

jacobhinkle · 2024-12-12T14:58:24Z

csrc/scheduler/hopper_multi_matmul.cpp

+  // Apply mma common transformation
+  c->split(-2, getM(params_->mma_macro));
+  c->split(-1, getN(params_->mma_macro));
+  // [..., Mo, No, Mio, Mii, Nio, Nii]
+  // -> [..., Mo, No, Mio, Nio, Mii, Nii]
+  c->reorder({{-3, -2}});
+  c->merge(-4);


I see this pattern a lot. Can we make this a utility method in this class and give it an informative name?

jacobhinkle · 2024-12-12T15:04:45Z

csrc/scheduler/hopper_multi_matmul.cpp

+  c->reorder({{-3, -2}});
+  c->merge(-4);
+
+  // [...., Mii, Nii] ->


This comment is a little hard for me to follow since there are some intermediate splits that are not shown. For example we start with Mii then get Miii(8) and Mii/16, so is that an (Mii/2)/8 (I think so)?
Since this is a complicated function maybe it is best to just skip to the last line (line 439) that shows the final sizes, and maybe give them different names.

jacobhinkle · 2024-12-12T15:05:58Z

csrc/scheduler/hopper_multi_matmul.cpp

+  auto s =
+      mma_utils::MmaSwizzler::scheduleMmaOutputAllocation(c->getLoopDomain());
+  c->setLoopDomain(s.as<IterDomain*>());
+  c->axis(-5)->parallelize(ParallelType::TIDy);


I mentioned in #3575 (comment) that this whole chunk should basically be a new method.

I'll wait for @rdspring1 to land his PR then use his fix or add one.

tests/cpp/test_matmul_scheduler.cpp

rdspring1

I'll wait to review until you pull the latests changes but the PR looks good to me. 🎉

protonu · 2024-12-13T18:47:06Z

!build

rdspring1 · 2024-12-13T18:51:47Z

csrc/scheduler/hopper_multi_matmul.cpp

-        2,
-        {c},
-        {ParallelType::BIDx, ParallelType::BIDy, ParallelType::BIDz});
+  // Propato to (not including) the splitk output if there is a splitk


Suggested change

// Propato to (not including) the splitk output if there is a splitk

// Propagate to (not including) the splitk output if there is a splitk

rdspring1

LGTM.

Just be aware of #3590. The current use_smem_epilogue epilogue will have conflicts. Your changes should be fine.

csrc/scheduler/hopper_multi_matmul.cpp

tests/cpp/test_matmul_scheduler.cpp

csrc/scheduler/hopper_multi_matmul.cpp

jacobhinkle · 2024-12-13T19:42:19Z

csrc/scheduler/hopper_multi_matmul.cpp

      scheduler_utils::BoundedDirectionalTransformPropagator::backward(
          d,
          -1,
-          mma_results_,
+          propagate_to,
          scheduler_utils::BoundedDirectionalTransformPropagator::Options()


Just a note: we'll need to revisit this path for cases where TMA is disabled. We need to make sure it is doing coalesced stores here (can check with ncu). I imagine that we will almost always want to use TMA, but it's possible that we would skip it if it increases smem usage too much. In those cases we'll hit this path.

csrc/scheduler/hopper_multi_matmul.cpp

jacobhinkle

LGTM. I think this gets us close to full functionality. BTW since smem epilogue is supported please change the title and PR description.

csrc/scheduler/hopper_multi_matmul.cpp

protonu · 2024-12-13T20:08:55Z

LGTM. I think this gets us close to full functionality. BTW since smem epilogue is supported please change the title and PR description.

Not as yet - this PR doesn't support smem epilogue.

jacobhinkle · 2024-12-13T20:24:02Z

LGTM. I think this gets us close to full functionality. BTW since smem epilogue is supported please change the title and PR description.

Not as yet - this PR doesn't support smem epilogue.

Oh OK nevermind then.

protonu · 2024-12-14T04:13:22Z

!test

protonu added the Matmuls label Dec 12, 2024

protonu requested review from zasdfgbnm, jacobhinkle and rdspring1 December 12, 2024 14:49

protonu commented Dec 12, 2024

View reviewed changes

tests/cpp/test_matmul_scheduler.cpp Show resolved Hide resolved

jacobhinkle approved these changes Dec 12, 2024

View reviewed changes

protonu mentioned this pull request Dec 12, 2024

Implement basic split-k gemm for hopper matmul scheduler #3575

Merged

rdspring1 reviewed Dec 13, 2024

View reviewed changes

protonu force-pushed the pbasu_mma_epilogue_hopper_no_smem_epilogue branch from 94f39a1 to be5f18b Compare December 13, 2024 18:41

protonu requested review from jacobhinkle and rdspring1 December 13, 2024 18:41

rdspring1 reviewed Dec 13, 2024

View reviewed changes

rdspring1 approved these changes Dec 13, 2024

View reviewed changes

jacobhinkle reviewed Dec 13, 2024

View reviewed changes

csrc/scheduler/hopper_multi_matmul.cpp Outdated Show resolved Hide resolved

csrc/scheduler/hopper_multi_matmul.cpp Show resolved Hide resolved

jacobhinkle requested changes Dec 13, 2024

View reviewed changes

protonu requested review from rdspring1 and jacobhinkle December 13, 2024 19:56

jacobhinkle approved these changes Dec 13, 2024

View reviewed changes

csrc/scheduler/hopper_multi_matmul.cpp Outdated Show resolved Hide resolved

protonu added 5 commits December 13, 2024 17:36

adding a new unit test for mma+bias and propating schedules

727ec7f

rebase and address reviewer comments

03180d1

removing header

34ec5e1

addressing reviewer comments

d87495b

fixing a bug

09cb407

protonu force-pushed the pbasu_mma_epilogue_hopper_no_smem_epilogue branch from 895df0f to 09cb407 Compare December 14, 2024 01:42

Merge branch 'main' into pbasu_mma_epilogue_hopper_no_smem_epilogue

dafd0c8

protonu merged commit d53be45 into main Dec 14, 2024
48 checks passed

protonu deleted the pbasu_mma_epilogue_hopper_no_smem_epilogue branch December 14, 2024 14:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Schedule epilogue (for Hopper Matmul) by propagation backward from output - smem epilogue not supported. #3580

Schedule epilogue (for Hopper Matmul) by propagation backward from output - smem epilogue not supported. #3580

protonu commented Dec 12, 2024 •

edited

Loading

jacobhinkle left a comment

jacobhinkle Dec 12, 2024

jacobhinkle Dec 12, 2024

jacobhinkle Dec 12, 2024

protonu Dec 12, 2024

rdspring1 left a comment

protonu commented Dec 13, 2024

rdspring1 Dec 13, 2024

rdspring1 left a comment

jacobhinkle Dec 13, 2024

jacobhinkle left a comment •

edited

Loading

protonu commented Dec 13, 2024

jacobhinkle commented Dec 13, 2024

protonu commented Dec 14, 2024

	// Propato to (not including) the splitk output if there is a splitk
	// Propagate to (not including) the splitk output if there is a splitk

Schedule epilogue (for Hopper Matmul) by propagation backward from output - smem epilogue not supported. #3580

Schedule epilogue (for Hopper Matmul) by propagation backward from output - smem epilogue not supported. #3580

Conversation

protonu commented Dec 12, 2024 • edited Loading

jacobhinkle left a comment

Choose a reason for hiding this comment

jacobhinkle Dec 12, 2024

Choose a reason for hiding this comment

jacobhinkle Dec 12, 2024

Choose a reason for hiding this comment

jacobhinkle Dec 12, 2024

Choose a reason for hiding this comment

protonu Dec 12, 2024

Choose a reason for hiding this comment

rdspring1 left a comment

Choose a reason for hiding this comment

protonu commented Dec 13, 2024

rdspring1 Dec 13, 2024

Choose a reason for hiding this comment

rdspring1 left a comment

Choose a reason for hiding this comment

jacobhinkle Dec 13, 2024

Choose a reason for hiding this comment

jacobhinkle left a comment • edited Loading

Choose a reason for hiding this comment

protonu commented Dec 13, 2024

jacobhinkle commented Dec 13, 2024

protonu commented Dec 14, 2024

protonu commented Dec 12, 2024 •

edited

Loading

jacobhinkle left a comment •

edited

Loading