MMA-izing the prolongator and restrictor kernels #1497

hummingtree · 2024-09-27T13:57:15Z

As the name suggests, this PR adds initial support for MMA-izing the prolongator and restrictor kernels. In addition,

Initial supports are also added to use tensor memory accelerator (TMA) for the memory movement on the Hopper and Blackwell GPUs;
Some general shard memory load/store patterns are optimized.

Remaining to-dos:

Add a command line interface for toggling MMA-ized prolongator and restrictor kernels.
Make the prolongator and restrictor working for staggered (the spin/parity handling is different).
A comprehensive testing of the MG/MMA workflows.
Add doxygen.
Clang-format

- Still need to add the spin factor of 2; - Still need to cover the to_non_rel = true;

…for loading from gmem.

…eature/prolongator-mma

… the process add an additional template to `kernel_param`.

…o feature/prolongator-mma

…underlying code such that FP16 works with rescaling.

…fields.

…eature/prolongator-mma

…ator_mma.in.cu` and `lib/restrictor.in.cpp` under the `QUDA_MMA_AVAILABLE` macro.

- Fix the issue with modifying the kenrel argument in fused MDW kernels; - Fix the issue where QUDA_ENABLE_MMA not being an option when MULTIGRID is disabled; - Fix compilation when MMA is not enabled.

…n 256.

…a_helper.hpp`.

include/color_spinor_field_order.h

include/targets/cuda/mma_tensor_op/smma_m16n8k8_sm70.cuh

include/targets/cuda/tma_helper.hpp

lib/block_transpose.in.cu

lib/prolongator.in.cpp

…pp`.

include/expand_list.hpp

include/kernels/dslash_coarse_mma.cuh

weinbe2 · 2024-11-12T17:54:50Z

include/targets/cuda/kernel.h


-    while (i < arg.threads.x) {
+    if constexpr (Arg::check_bounds) {


Should this go into, for ex, the hip target for completeness/consistency even if it's unused for now?

include/targets/cuda/mma_tensor_op/gmem_loader.cuh

lib/prolongator_mma.in.cu

…indices).

hummingtree and others added 28 commits August 27, 2024 12:20

Add MMA version of prolongator.

88a108e

Dagger'ed the equation to make better use of MMA.

c0cb6e0

Make nColor = 3 works:

dd55d92

- Still need to add the spin factor of 2; - Still need to cover the to_non_rel = true;

Add from_to_non_rel to block transpose to complete the circle.

57b6f95

More cleanup of the MMA code. Apply vector gmem loads when possible.

efa9c15

Apply more vector gmem loads when possible.

82d532b

Add MMA version for restrictor.

e6a5d34

Add expands for restrictor with MMA.

f0f88d5

Add shared memory caching for the the restrictor kernel.

92b1c39

Allow restrictor_mma to have N % bN != 0; More generic optimizations …

e7b9361

…for loading from gmem.

Add rescaling to prolongator.

6bb02f6

Modify the MMA types.

2d8867e

Add rescale for restrictor.

db6c030

Abstract the MMA expansions into a class.

ce257ef

Add more abstraction; make TF32 the default for SM80 and later.

f38882a

Merge branch 'feature/mrhs-solvers' of github.com:lattice/quda into f…

69e15d6

…eature/prolongator-mma

Merge branch 'feature/mrhs-solvers' of github.com:lattice/quda into f…

e03ca23

…eature/prolongator-mma

Fix block transpose by having no bound checks on the kernel level; in…

0a2bd2a

… the process add an additional template to `kernel_param`.

Merge branch 'feature/prolongator-mma' of github.com:lattice/quda int…

e68dfa9

…o feature/prolongator-mma

Add some aggregate sizes to MMA restrictor

de05ebf

Make aggregate_size a runtime variable.

bcdfb86

Apply clang-format.

189f78e

Soften the restriction for nrhs from multiple of 16 to multiple of 8.

8636160

Set the default precision in coarse dslash mma to TF32/FP16; Fix the …

f3a42bd

…underlying code such that FP16 works with rescaling.

Clean up the MMA code.

82a61c6

Short cut the rescaling code to use scale_inv for fixed point format …

6515177

…fields.

Clean up code.

a9f6d7c

Merge branch 'feature/mrhs-solvers' of github.com:lattice/quda into f…

29593b1

…eature/prolongator-mma

hummingtree requested review from a team as code owners September 27, 2024 13:57

hummingtree added 12 commits October 31, 2024 17:18

Put tma_wait under the TMA macro.

c3a155d

Add const to include/kernels/dslash_mdw_fused.cuh; Put `lib/prolong…

88bedc9

…ator_mma.in.cu` and `lib/restrictor.in.cpp` under the `QUDA_MMA_AVAILABLE` macro.

Add more const to Arg &arg; Fix compile when MMA is not available.

013aa50

Fix CI:

48af0e2

- Fix the issue with modifying the kenrel argument in fused MDW kernels; - Fix the issue where QUDA_ENABLE_MMA not being an option when MULTIGRID is disabled; - Fix compilation when MMA is not enabled.

Remove the unwanted warning macro.

123aa04

Address the unused parameter warnings from CI.

c6e833f

Add support for staggered nSpin = 1 case for restrictor_mma.

bb83b5f

Break TMA box sizes into smaller pieces when the sizes are larger tha…

7f31418

…n 256.

Restore the SMMA precision for coarse dslash MMA.

3a0f205

Use SIMT for coarse dslash MMA on SM70.

6d6beb5

Add param and command line option for running transfer with MMA.

43e7cf7

Add doxygen to include/expand_list.hpp and `include/targets/cuda/tm…

53eb2ae

…a_helper.hpp`.

weinbe2 reviewed Nov 7, 2024

View reviewed changes

include/color_spinor_field_order.h Outdated Show resolved Hide resolved

weinbe2 reviewed Nov 8, 2024

View reviewed changes

include/targets/cuda/mma_tensor_op/smma_m16n8k8_sm70.cuh Outdated Show resolved Hide resolved

weinbe2 reviewed Nov 8, 2024

View reviewed changes

include/targets/cuda/tma_helper.hpp Outdated Show resolved Hide resolved

weinbe2 reviewed Nov 8, 2024

View reviewed changes

lib/block_transpose.in.cu Show resolved Hide resolved

weinbe2 reviewed Nov 8, 2024

View reviewed changes

lib/prolongator.in.cpp Outdated Show resolved Hide resolved

hummingtree added 3 commits November 12, 2024 06:44

Fix typo; replace store_type with store_t.

d0da823

Only instantiate enabled spins in lib/block_transpose.in.cu.

cd3bbad

Remove test code in lib/restrictor.in.cpp and `lib/prolongator.in.c…

1be04f5

…pp`.

weinbe2 reviewed Nov 12, 2024

View reviewed changes

include/expand_list.hpp Outdated Show resolved Hide resolved

weinbe2 reviewed Nov 12, 2024

View reviewed changes

include/kernels/dslash_coarse_mma.cuh Show resolved Hide resolved

weinbe2 reviewed Nov 12, 2024

View reviewed changes

include/targets/cuda/mma_tensor_op/gmem_loader.cuh Outdated Show resolved Hide resolved

weinbe2 reviewed Nov 12, 2024

View reviewed changes

lib/prolongator_mma.in.cu Show resolved Hide resolved

hummingtree added 5 commits November 12, 2024 20:20

Store the aux values as the shape values directly (instead of as the …

c84c376

…indices).

Use mma::numeric_limits<mma::half>::max() in MMA code.

6451e7f

Add initializer to *_factors.

eacc613

Use curly brace initializers.

24c925f

Initialize looping variable.

fd4a930

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

MMA-izing the prolongator and restrictor kernels #1497

MMA-izing the prolongator and restrictor kernels #1497

hummingtree commented Sep 27, 2024 •

edited

Loading

weinbe2 Nov 12, 2024


		while (i < arg.threads.x) {
		if constexpr (Arg::check_bounds) {

MMA-izing the prolongator and restrictor kernels #1497

Are you sure you want to change the base?

MMA-izing the prolongator and restrictor kernels #1497

Conversation

hummingtree commented Sep 27, 2024 • edited Loading

weinbe2 Nov 12, 2024

Choose a reason for hiding this comment

hummingtree commented Sep 27, 2024 •

edited

Loading