dmma_tma_kernel.h

Go to the documentation of this file.

#ifndef __MATH_DMMA_TMA_KERNEL_H__
#define __MATH_DMMA_TMA_KERNEL_H__
/*
 Copyright (c) 2026, The Neko Authors
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions
 are met:
 
   * Redistributions of source code must retain the above copyright
     notice, this list of conditions and the following disclaimer.
 
   * Redistributions in binary form must reproduce the above
     copyright notice, this list of conditions and the following
     disclaimer in the documentation and/or other materials provided
     with the distribution.
 
   * Neither the name of the authors nor the names of its
     contributors may be used to endorse or promote products derived
     from this software without specific prior written permission.
 
 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 POSSIBILITY OF SUCH DAMAGE.
*/
 
#include <stdlib.h>
#include <stdint.h>
#include <stddef.h>
#include <cuda_runtime.h>
#include <device/device_config.h>
#include "dmma_kernel.h"
 
/* The geometric factor cubes staged per element: h1, g11, g22, g33, g12,
   g13, g23, in that order */
#define DMMA_NG 7
 
/* opgrad stages the nine metric cubes drdx..dtdz. Its tenth factor, the
   quadrature weight w3, is shared by every element and so is read straight
   from global memory -- it is L2 resident after the first block and a bulk
   copy of it per element would be pure waste */
#define DMMA_NG_OPGRAD 9
 
/* conv1 stages the same nine, plus the three convecting velocity components
   and jacinv */
#define DMMA_NG_CONV1 13
 
/*
 * cp.async.bulk is PTX ISA 8.0, which means CUDA 12.0. Compiling for sm_90 is
 * not enough on its own: CUDA 11.8 targets sm_90 but its ptxas only speaks
 * PTX 7.8 and would reject every asm block below, so the toolkit version is
 * part of the guard rather than inferred from the arch.
 */
#if defined(CUDART_VERSION) && (CUDART_VERSION >= 12000)
#define NEKO_TMA_TOOLKIT 1
#else
#define NEKO_TMA_TOOLKIT 0
#endif
 
/*
 * Was sm_90 among the architectures this translation unit was compiled for?
 * Same trap as dmma_arch_compiled() in dmma_kernel.h: a binary built for an
 * older arch still runs on Hopper by JIT compiling its PTX, and the kernel
 * body is guarded on __CUDA_ARCH__, so without this check a build for, say,
 * sm_80 running on a GH200 would JIT the *no-op* body, time as free, win the
 * tuner comparison and leave a stale w. Unlike the DMMA check this one is
 * sm_90 exactly, not a range: sm_100 and later move the bulk copy semantics
 * and would have to be measured before being allowed in.
 */
#ifndef NEKO_TMA_ARCH_COMPILED
#if defined(__CUDA_ARCH_LIST__)
static inline bool tma_arch_compiled()
{
  const int arch[] = { __CUDA_ARCH_LIST__ };
 
  for (int i = 0; i < (int) (sizeof(arch)/sizeof(arch[0])); i++) {
    if (arch[i] >= 900 && arch[i] < 1000) {
      return true;
    }
  }
  return false;
}
#else
static inline bool tma_arch_compiled()
{
  return false;
}
#endif
#else
static inline bool tma_arch_compiled()
{
  return (NEKO_TMA_ARCH_COMPILED != 0);
}
#endif
 
static inline bool cuda_have_tma()
{
#if NEKO_TMA_TOOLKIT
  static int cached = -1;
 
  if (cached < 0) {
    int dev = 0;
    cudaDeviceProp prop;
 
    if (!tma_arch_compiled()) {
      cached = 0;
    } else if (cudaGetDevice(&dev) == cudaSuccess &&
               cudaGetDeviceProperties(&prop, dev) == cudaSuccess) {
      cached = (prop.major == 9) ? 1 : 0;
    } else {
      cached = 0;
    }
  }
  return cached == 1;
#else
  return false;
#endif
}
 
template< const int LX >
static inline bool dmma_tma_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
template< const int LX >
static inline bool dmma_tma_vector_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
template< const int LX >
static inline bool dmma_tma_dudxyz_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
template< const int LX >
static inline bool dmma_tma_opgrad_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
template< const int LX >
static inline bool dmma_tma_conv1_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
template< const int LX >
static inline bool dmma_tma_cdtp_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
/*
 * The batched vector variant's block, as one struct in dynamic shared memory.
 *
 * Why a struct and why dynamic. The measurement that motivates this variant is
 * that a bulk copy only pays for itself in company: the scalar kernel issues
 * eight at once and wins at 20 warps per SM against 48, while the
 * component-at-a-time vector kernel issues seven once and then two lone 4 kB
 * copies, and loses at 20 warps against 16 -- with MORE occupancy than the
 * kernel it loses to, which is what rules residency out as the cause. Putting
 * every one of an element's thirteen cubes in shared at once takes the block
 * to 54800 B, past the 48 kB a block gets without opting in, so the allocation
 * has to be dynamic and the layout has to live somewhere. A struct puts it in
 * one place and makes the launch size sizeof(), which cannot drift from it.
 *
 * Member order is the layout: every cube lands on a 4 kB boundary and the two
 * mbarriers go last, so no padding is inserted and every bulk copy target is
 * aligned far past the 16 bytes it needs. The asserts below hold that.
 */
struct dmma_tma_batch_smem {
  double c[3][DMMA_CUBE];        /* the three components, in and then out */
  double r[DMMA_CUBE];           /* reference derivatives, reused by each */
  double s[DMMA_CUBE];           /* component in turn */
  double t[DMMA_CUBE];
  double g[DMMA_NG][DMMA_CUBE];  /* h1, g11, g22, g33, g12, g13, g23 */
  double dx[DMMA_MAT];
  double dy[DMMA_MAT];
  double dz[DMMA_MAT];
  unsigned long long bar_a;      /* the first component's cube */
  unsigned long long bar_b;      /* the other two, and the seven factors */
};
 
#define NEKO_DMMA_TMA_BATCH_SMEM ((int) sizeof(dmma_tma_batch_smem))
 
/* 227 kB is the most a Hopper block can opt into; the runtime gate below
   checks the device rather than trusting this */
static_assert(sizeof(dmma_tma_batch_smem) <= 227 * 1024,
              "dmma tma batch block exceeds the opt-in shared memory ceiling");
static_assert(offsetof(dmma_tma_batch_smem, c) % 128 == 0 &&
              offsetof(dmma_tma_batch_smem, g) % 128 == 0,
              "dmma tma batch bulk copy targets must stay 128 byte aligned");
 
static inline int cuda_tma_smem_optin()
{
#if NEKO_TMA_TOOLKIT
  static int cached = -1;
 
  if (cached < 0) {
    int dev = 0;
    int optin = 0;
 
    if (!cuda_have_tma()) {
      cached = 0;
    } else if (cudaGetDevice(&dev) == cudaSuccess &&
               cudaDeviceGetAttribute(&optin,
                                      cudaDevAttrMaxSharedMemoryPerBlockOptin,
                                      dev) == cudaSuccess) {
      cached = optin;
    } else {
      cached = 0;
    }
  }
  return cached;
#else
  return 0;
#endif
}
 
static inline bool cuda_have_tma_batch()
{
  return cuda_tma_smem_optin() >= NEKO_DMMA_TMA_BATCH_SMEM;
}
 
/*
 * opgrad's block, laid out the same way and for the same reason as
 * dmma_tma_batch_smem above: thirteen cubes will not fit in the 48 kB a block
 * gets without asking, so the allocation is dynamic and the layout lives in a
 * struct whose sizeof() is the launch size.
 *
 * It comes to **exactly** the same 54800 B as the batched axhelm block -- both
 * are thirteen cubes and three matrices -- so cuda_have_tma_batch()'s device
 * gate covers this variant unchanged, and so does its four blocks per SM.
 *
 * The ten input copies are issued together, which is the arrangement the
 * scalar axhelm variant won with and the component-at-a-time vector one lost
 * with. The three outputs leave as ordinary coalesced stores rather than bulk
 * ones: only 'u' is free by then, so bulk storing all three would either
 * serialise on a single cube or cost a register hoist of rtmp/stmp/ttmp --
 * and at 19% of the traffic the stores are not what this variant is for.
 */
struct opgrad_tma_smem {
  double u[DMMA_CUBE];                  /* the input, dead after phase 1 */
  double r[DMMA_CUBE];                  /* the reference derivatives */
  double s[DMMA_CUBE];
  double t[DMMA_CUBE];
  double g[DMMA_NG_OPGRAD][DMMA_CUBE];  /* drdx, dsdx, dtdx, drdy, ... dtdz */
  double dx[DMMA_MAT];
  double dy[DMMA_MAT];
  double dz[DMMA_MAT];
  unsigned long long bar_u;             /* the input cube */
  unsigned long long bar_g;             /* the nine metric cubes */
};
 
#define NEKO_OPGRAD_TMA_SMEM ((int) sizeof(opgrad_tma_smem))
 
/*
 * conv1's block. Four more cubes than opgrad -- the three convecting velocity
 * components and jacinv -- because every one of them is consumed at the same
 * pointwise step and batching all fourteen copies is the whole point.
 *
 * 71184 B is three blocks per SM rather than opgrad's four, which is below the
 * batched axhelm kernel's four and is the reason this one is a genuine
 * question rather than an expected win. It has the best bytes per contraction
 * of any operator in the tree (20480 at lx = 8, against axhelm's 6144), so
 * there is more memory time here to hide the staging behind than anywhere
 * else; whether that beats losing a block is what the tuner is for. Unlike
 * opgrad there is a single output, so 'u' is free for it and the result does
 * leave as one bulk store.
 */
struct conv1_tma_smem {
  double u[DMMA_CUBE];                  /* the input, then the output */
  double r[DMMA_CUBE];
  double s[DMMA_CUBE];
  double t[DMMA_CUBE];
  double g[DMMA_NG_CONV1][DMMA_CUBE];   /* vx, vy, vz, jacinv, drdx .. dtdz */
  double dx[DMMA_MAT];
  double dy[DMMA_MAT];
  double dz[DMMA_MAT];
  unsigned long long bar_u;
  unsigned long long bar_g;
};
 
#define NEKO_CONV1_TMA_SMEM ((int) sizeof(conv1_tma_smem))
 
static_assert(sizeof(opgrad_tma_smem) <= 227 * 1024 &&
              sizeof(conv1_tma_smem) <= 227 * 1024,
              "tma block exceeds the opt-in shared memory ceiling");
static_assert(offsetof(opgrad_tma_smem, u) % 128 == 0 &&
              offsetof(opgrad_tma_smem, g) % 128 == 0 &&
              offsetof(conv1_tma_smem, u) % 128 == 0 &&
              offsetof(conv1_tma_smem, g) % 128 == 0,
              "tma bulk copy targets must stay 128 byte aligned");
 
static inline bool cuda_have_tma_opgrad()
{
  return cuda_tma_smem_optin() >= NEKO_OPGRAD_TMA_SMEM;
}
 
static inline bool cuda_have_tma_conv1()
{
  return cuda_tma_smem_optin() >= NEKO_CONV1_TMA_SMEM;
}
 
template< const int LX >
static inline bool dmma_tma_batch_lx_supported()
{
  return (sizeof(real) == 8) && (LX == DMMA_P);
}
 
static inline bool dmma_tma_ptr_aligned(const void *p)
{
  return ((((uintptr_t) p) & (uintptr_t) 15) == 0);
}
 
static inline bool dmma_tma_aligned(const void *w, const void *u,
                                    const void *h1, const void *g11,
                                    const void *g22, const void *g33,
                                    const void *g12, const void *g13,
                                    const void *g23)
{
  return dmma_tma_ptr_aligned(w) && dmma_tma_ptr_aligned(u) &&
         dmma_tma_ptr_aligned(h1) && dmma_tma_ptr_aligned(g11) &&
         dmma_tma_ptr_aligned(g22) && dmma_tma_ptr_aligned(g33) &&
         dmma_tma_ptr_aligned(g12) && dmma_tma_ptr_aligned(g13) &&
         dmma_tma_ptr_aligned(g23);
}
 
/* The vector operator's thirteen arrays, same check */
static inline bool dmma_tma_vector_aligned(const void *au, const void *av,
                                           const void *aw, const void *u,
                                           const void *v, const void *w,
                                           const void *h1, const void *g11,
                                           const void *g22, const void *g33,
                                           const void *g12, const void *g13,
                                           const void *g23)
{
  return dmma_tma_aligned(au, u, h1, g11, g22, g33, g12, g13, g23) &&
         dmma_tma_ptr_aligned(av) && dmma_tma_ptr_aligned(aw) &&
         dmma_tma_ptr_aligned(v) && dmma_tma_ptr_aligned(w);
}
 
/* The derivative operator's six arrays, same check */
static inline bool dmma_tma_dudxyz_aligned(const void *du, const void *u,
                                           const void *dr, const void *ds,
                                           const void *dt, const void *jacinv)
{
  return dmma_tma_ptr_aligned(du) && dmma_tma_ptr_aligned(u) &&
         dmma_tma_ptr_aligned(dr) && dmma_tma_ptr_aligned(ds) &&
         dmma_tma_ptr_aligned(dt) && dmma_tma_ptr_aligned(jacinv);
}
 
/* The nine metric cubes, which both opgrad and conv1 bulk copy */
static inline bool dmma_tma_metrics_aligned(const void *drdx, const void *dsdx,
                                            const void *dtdx, const void *drdy,
                                            const void *dsdy, const void *dtdy,
                                            const void *drdz, const void *dsdz,
                                            const void *dtdz)
{
  return dmma_tma_ptr_aligned(drdx) && dmma_tma_ptr_aligned(dsdx) &&
         dmma_tma_ptr_aligned(dtdx) && dmma_tma_ptr_aligned(drdy) &&
         dmma_tma_ptr_aligned(dsdy) && dmma_tma_ptr_aligned(dtdy) &&
         dmma_tma_ptr_aligned(drdz) && dmma_tma_ptr_aligned(dsdz) &&
         dmma_tma_ptr_aligned(dtdz);
}
 
/*
 * opgrad copies ten arrays in and stores three with ordinary stores, so only
 * the ten are checked -- an ordinary store has no alignment requirement beyond
 * the type's, and w3 is never bulk copied either.
 */
static inline bool dmma_tma_opgrad_aligned(const void *u,
                                           const void *drdx, const void *dsdx,
                                           const void *dtdx, const void *drdy,
                                           const void *dsdy, const void *dtdy,
                                           const void *drdz, const void *dsdz,
                                           const void *dtdz)
{
  return dmma_tma_ptr_aligned(u) &&
         dmma_tma_metrics_aligned(drdx, dsdx, dtdx, drdy, dsdy, dtdy,
                                  drdz, dsdz, dtdz);
}
 
/*
 * cdtp copies four in and bulk stores its single output. w3 is shared by every
 * element and read straight from global, so it is not among them.
 */
static inline bool dmma_tma_cdtp_aligned(const void *dtx, const void *x,
                                         const void *dr, const void *ds,
                                         const void *dt)
{
  return dmma_tma_ptr_aligned(dtx) && dmma_tma_ptr_aligned(x) &&
         dmma_tma_ptr_aligned(dr) && dmma_tma_ptr_aligned(ds) &&
         dmma_tma_ptr_aligned(dt);
}
 
/* conv1 copies fourteen in and bulk stores its single output, so all fifteen */
static inline bool dmma_tma_conv1_aligned(const void *du, const void *u,
                                          const void *vx, const void *vy,
                                          const void *vz, const void *jacinv,
                                          const void *drdx, const void *dsdx,
                                          const void *dtdx, const void *drdy,
                                          const void *dsdy, const void *dtdy,
                                          const void *drdz, const void *dsdz,
                                          const void *dtdz)
{
  return dmma_tma_ptr_aligned(du) && dmma_tma_ptr_aligned(u) &&
         dmma_tma_ptr_aligned(vx) && dmma_tma_ptr_aligned(vy) &&
         dmma_tma_ptr_aligned(vz) && dmma_tma_ptr_aligned(jacinv) &&
         dmma_tma_metrics_aligned(drdx, dsdx, dtdx, drdy, dsdy, dtdy,
                                  drdz, dsdz, dtdz);
}
 
/* Forced candidate, used when NEKO_AUTOTUNE pins the DMMA_TMA variant. The
   warps per block candidates are the DMMA ones, see NEKO_DMMA_NW() */
static int neko_dmma_tma_env()
{
  const char *v = getenv("NEKO_DMMA_TMA_NW");
  int c = (v != NULL) ? atoi(v) : 0;
 
  if (c < 0 || c >= NEKO_DMMA_CANDIDATES) {
    c = 0;
  }
  return c;
}
 
/* Report every measured DMMA_TMA_BATCH candidate, see NEKO_TUNE_LOG_DMMA */
#define NEKO_TUNE_LOG_DMMA_TMA_BATCH(LX, T5)                                  \
  do {                                                                        \
    for (int c = 0; c < NEKO_DMMA_CANDIDATES; c++) {                          \
      if ((T5)[c] >= NEKO_TUNE_INIT) { continue; }                            \
      char lbl_[16];                                                          \
      sprintf(lbl_, "TMAB  %dw", NEKO_DMMA_NW(c));                            \
      sprintf(neko_log_buf, "%-13s: %9.2f us/call", lbl_,                     \
              NEKO_TUNE_US((T5)[c], iters));                                  \
      log_message(neko_log_buf);                                              \
    }                                                                         \
  } while (0)
 
/* Report every measured DMMA_TMA candidate, see NEKO_TUNE_LOG_DMMA */
#define NEKO_TUNE_LOG_DMMA_TMA(LX, T4)                                        \
  do {                                                                        \
    for (int c = 0; c < NEKO_DMMA_CANDIDATES; c++) {                          \
      if ((T4)[c] >= NEKO_TUNE_INIT) { continue; }                            \
      char lbl_[16];                                                          \
      sprintf(lbl_, "TMA   %dw", NEKO_DMMA_NW(c));                            \
      sprintf(neko_log_buf, "%-13s: %9.2f us/call", lbl_,                     \
              NEKO_TUNE_US((T4)[c], iters));                                  \
      log_message(neko_log_buf);                                              \
    }                                                                         \
  } while (0)
 
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) &&                       \
    (__CUDA_ARCH__ < 1000) && NEKO_TMA_TOOLKIT
 
/*
 * The mbarrier and bulk copy primitives, as inline PTX.
 *
 * Written out rather than reached through cuda::memcpy_async() and
 * cuda::barrier because the point of the variant is to measure the TMA
 * engine: libcu++ selects cp.async.bulk only when it can prove the size and
 * alignment, and falls back to cp.async otherwise, which would still compile,
 * still be correct, and quietly measure something else entirely.
 *
 * Every mbarrier and shared operand is addressed in the shared window, which
 * is what __cvta_generic_to_shared() produces.
 */
 
__device__ __forceinline__ static unsigned int tma_smem(const void *p)
{
  return (unsigned int) __cvta_generic_to_shared(p);
}
 
/* And the global window, likewise. The bulk copies name the global state
   space explicitly, so a generic address is not theirs to reinterpret: the
   conversion is written out rather than left to the pointer happening to be
   global already */
__device__ __forceinline__ static unsigned long long tma_gmem(const void *p)
{
  return (unsigned long long) __cvta_generic_to_global(p);
}
 
/* 'count' is the number of arrivals the barrier waits for, one per producer
   thread. Must be visible to every waiter before any of them waits, i.e. a
   __syncthreads() has to separate this from the first tma_wait() */
__device__ __forceinline__ static void tma_barrier_init(unsigned long long *bar,
                                                        const int count)
{
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;"
               :: "r"(tma_smem(bar)), "r"(count) : "memory");
}
 
/* Arrive, and declare how many bytes the copies about to be issued will
   deliver. The barrier flips phase once the arrival has landed *and* the
   completing copies have paid off the transaction count, so this has to
   precede the copies it accounts for */
__device__ __forceinline__ static void tma_expect(unsigned long long *bar,
                                                  const int bytes)
{
  asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;"
               :: "r"(tma_smem(bar)), "r"(bytes) : "memory");
}
 
/* One contiguous global -> shared bulk copy, completing on 'bar' */
__device__ __forceinline__ static void tma_load(void *dst, const void *src,
                                                const int bytes,
                                                unsigned long long *bar)
{
  asm volatile("cp.async.bulk.shared::cluster.global"
               ".mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];"
               :: "r"(tma_smem(dst)), "l"(tma_gmem(src)), "r"(bytes),
                  "r"(tma_smem(bar)) : "memory");
}
 
/* Spin until 'bar' completes the phase of parity 'phase'; the initial parity
   is 0. try_wait in a C loop rather than a wait loop inside the asm because
   this is inlined more than once per kernel and PTX labels are not renamed on
   inlining, so an asm local label would be defined twice.
 
   The volatile asm and its memory clobber are also what keep the compiler
   from sinking the wait below the shared reads it guards; the hardware
   ordering is the mbarrier's own. */
__device__ __forceinline__ static void tma_wait(unsigned long long *bar,
                                                const int phase)
{
  unsigned int done;
 
  do {
    asm volatile("{\n"
                 ".reg .pred p;\n"
                 "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n"
                 "selp.b32 %0, 1, 0, p;\n"
                 "}\n"
                 : "=r"(done) : "r"(tma_smem(bar)), "r"(phase) : "memory");
  } while (done == 0);
}
 
/* Make this thread's ordinary shared memory writes visible to the async
   proxy. Reads of shared memory by a bulk copy go through that proxy, and
   ordinary stores are not ordered before them by __syncthreads() alone, so a
   store issued without this can read stale shared memory. The load direction
   needs no counterpart: there the mbarrier completion is what orders the
   copy's writes before the ordinary reads that consume them. */
__device__ __forceinline__ static void tma_fence_shared()
{
  asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
}
 
/* One contiguous shared -> global bulk copy. Tracked by a bulk group rather
   than an mbarrier: nothing reads the result back, it only has to have been
   committed before the block releases its shared memory */
__device__ __forceinline__ static void tma_store(void *dst, const void *src,
                                                 const int bytes)
{
  asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;"
               :: "l"(tma_gmem(dst)), "r"(tma_smem(src)), "r"(bytes)
               : "memory");
}
 
/* Commit the issued stores and wait for them to have read their source out of
   shared memory. Only the issuing thread has a group to wait on, but no
   thread may leave the block before it completes, hence the __syncthreads()
   that follows this at the call site */
__device__ __forceinline__ static void tma_store_wait()
{
  asm volatile("cp.async.bulk.commit_group;" ::: "memory");
  asm volatile("cp.async.bulk.wait_group 0;" ::: "memory");
}
 
#endif // __CUDA_ARCH__ == sm_90 with a CUDA 12 toolkit
 
#endif // __MATH_DMMA_TMA_KERNEL_H__