copy_sm90.h

#pragma once

#ifndef __CUDACC_RTC__
#include <cuda.h>
#endif

#include "common.h"

namespace tl {
enum class CacheHintSm90 : uint64_t {
  EVICT_NORMAL = 0x1000000000000000,
  EVICT_FIRST = 0x12F0000000000000,
  EVICT_LAST = 0x14F0000000000000,
};

TL_DEVICE void tma_load(void *smem_ptr, void *gmem_ptr, uint64_t &smem_mbar,
                        uint32_t size) {
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::"
               "bytes [%0], [%1], %2, [%3]; \n" ::"r"(smem_int_ptr),
               "l"(gmem_ptr), "r"(size), "r"(smem_int_mbar)
               :);
}

TL_DEVICE void tma_load_multicast(void *smem_ptr, void *gmem_ptr,
                                  uint64_t &smem_mbar, uint32_t size,
                                  uint16_t mask) {
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  asm volatile(
      "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes."
      "multicast::cluster [%0], [%1], %2, [%3], %4; \n" ::"r"(smem_int_ptr),
      "l"(gmem_ptr), "r"(size), "r"(smem_int_mbar), "h"(mask)
      :);
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
                        void const *const smem_ptr, int32_t const &crd0) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes"
                 " [%0], [%1, {%3}], [%2];"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes.L2::cache_hint"
                 " [%0], [%1, {%3}], [%2], %4;"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
                        void const *const smem_ptr, int32_t const &crd0,
                        int32_t const &crd1) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes"
                 " [%0], [%1, {%3, %4}], [%2];"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes.L2::cache_hint"
                 " [%0], [%1, {%3, %4}], [%2], %5;"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
                        void const *const smem_ptr, int32_t const &crd0,
                        int32_t const &crd1, int32_t const &crd2) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes"
                 " [%0], [%1, {%3, %4, %5}], [%2];"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "r"(crd2)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes.L2::cache_hint"
                 " [%0], [%1, {%3, %4, %5}], [%2], %6;"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "r"(crd2), "l"(cache_hint)
                 : "memory");
  }
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
                        void const *const smem_ptr, int32_t const &crd0,
                        int32_t const &crd1, int32_t const &crd2,
                        int32_t const &crd3) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes"
                 " [%0], [%1, {%3, %4, %5, %6}], [%2];"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes.L2::cache_hint"
                 " [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
                        void const *const smem_ptr, int32_t const &crd0,
                        int32_t const &crd1, int32_t const &crd2,
                        int32_t const &crd3, int32_t const &crd4) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes"
                 " [%0], [%1, {%3, %4, %5, %6, %7}], [%2];"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::"
                 "complete_tx::bytes.L2::cache_hint"
                 " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
                 :
                 : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
                   "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4),
                   "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load_im2col(const CUtensorMap &descriptor,
                               uint64_t &smem_mbar, void const *const smem_ptr,
                               int32_t const &coord_c, int32_t const &coord_w,
                               int32_t const &coord_h, int32_t const &coord_n,
                               uint16_t const &offset_w,
                               uint16_t const &offset_h) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile(
        "cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier:"
        ":complete_tx::bytes"
        " [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8};"
        :
        : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
          "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), "h"(offset_w),
          "h"(offset_h)
        : "memory");
  } else {
    asm volatile(
        "cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier:"
        ":complete_tx::bytes.L2::cache_hint"
        " [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8}, %9;"
        :
        : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
          "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), "h"(offset_w),
          "h"(offset_h), "l"(cache_hint)
        : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
                         void const *const smem_ptr, int32_t const &crd0) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);

  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.1d.global.shared::cta.bulk_group [%0, "
                 "{%2}], [%1];"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.1d.global.shared::cta.bulk_group "
                 "::cache_hint [%0, {%2}], [%1], %3;"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0),
                   "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
                         void const *const smem_ptr, int32_t const &crd0,
                         int32_t const &crd1) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);

  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, "
                 "{%2, %3}], [%1];"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.bulk_group "
                 "::cache_hint [%0, {%2, %3}], [%1], %4;"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
                         void const *const smem_ptr, int32_t const &crd0,
                         int32_t const &crd1, int32_t const &crd2) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);

  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.3d.global.shared::cta.bulk_group [%0, "
                 "{%2, %3, %4}], [%1];"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "r"(crd2)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.3d.global.shared::cta.bulk_group "
                 "::cache_hint [%0, {%2, %3, %4}], [%1], %5;"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "r"(crd2), "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
                         void const *const smem_ptr, int32_t const &crd0,
                         int32_t const &crd1, int32_t const &crd2,
                         int32_t const &crd3) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);

  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.4d.global.shared::cta.bulk_group [%0, "
                 "{%2, %3, %4, %5}], [%1];"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "r"(crd2), "r"(crd3)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.4d.global.shared::cta.bulk_group "
                 "::cache_hint [%0, {%2, %3, %4, %5}], [%1], %6;"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "r"(crd2), "r"(crd3), "l"(cache_hint)
                 : "memory");
  }
}

template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
                         void const *const smem_ptr, int32_t const &crd0,
                         int32_t const &crd1, int32_t const &crd2,
                         int32_t const &crd3, int32_t const &crd4) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);

  if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
    asm volatile("cp.async.bulk.tensor.5d.global.shared::cta.bulk_group [%0, "
                 "{%2, %3, %4, %5, %6}], [%1];"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "r"(crd2), "r"(crd3), "r"(crd4)
                 : "memory");
  } else {
    asm volatile("cp.async.bulk.tensor.5d.global.shared::cta.bulk_group "
                 "::cache_hint [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
                 :
                 : "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
                   "r"(crd2), "r"(crd3), "r"(crd4), "l"(cache_hint)
                 : "memory");
  }
}

TL_DEVICE void prefetch_tma_descriptor(const CUtensorMap &descriptor) {
  uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
  asm volatile("prefetch.tensormap [%0];" : : "l"(gmem_int_desc) : "memory");
}

TL_DEVICE void mbarrier_init(uint64_t &smem_barrier, uint32_t arrive_count) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  asm volatile("mbarrier.init.shared.b64 [%1], %0;"
               :
               : "r"(arrive_count), "r"(smem_int_ptr));
}

TL_DEVICE uint32_t mbarrier_try_wait(uint64_t &smem_barrier, int phase_bit) {

  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  uint32_t waitComplete;

  asm volatile("{\n\t"
               ".reg .pred P1; \n\t"
               "mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \n\t"
               "selp.b32 %0, 1, 0, P1; \n\t"
               "}"
               : "=r"(waitComplete)
               : "r"(smem_int_ptr), "r"(phase_bit));

  return waitComplete;
}

TL_DEVICE void mbarrier_wait(uint64_t &smem_barrier, int phase_bit) {
  if (mbarrier_try_wait(smem_barrier, phase_bit) == 0) {
    uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
    // Arbitrarily large timer value after which try-wait expires and re-tries.
    uint32_t ticks = 0x989680;
    asm volatile("{\n\t"
                 ".reg .pred       P1; \n\t"
                 "LAB_WAIT: \n\t"
                 "mbarrier.try_wait.parity.shared.b64 P1, [%0], %1, %2; \n\t"
                 "@P1 bra DONE; \n\t"
                 "bra     LAB_WAIT; \n\t"
                 "DONE: \n\t"
                 "}"
                 :
                 : "r"(smem_int_ptr), "r"(phase_bit), "r"(ticks));
  }
}

TL_DEVICE void mbarrier_test_wait(uint64_t &smem_barrier, int phase_bit) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  asm volatile(
      "{\n"
      ".reg .pred                P1;\n"
      "LAB_WAIT:\n"
      "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
      "@P1                       bra.uni DONE;\n"
      "nanosleep.u32 5;\n" // wait a few nanoseconds on pre-Hopper architectures
                           // to save instruction issue slots
      "bra.uni                   LAB_WAIT;\n"
      "DONE:\n"
      "}\n" ::"r"(smem_int_ptr),
      "r"(phase_bit));
}

TL_DEVICE void mbarrier_arrive(uint64_t &smem_barrier) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  asm volatile("mbarrier.arrive.shared.b64 _, [%0];" : : "r"(smem_int_ptr));
}

TL_DEVICE void mbarrier_arrive(uint64_t &smem_barrier, int cta_id,
                               uint32_t pred) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  if (pred) {
    asm volatile("{\n\t"
                 ".reg .b32 remAddr32;\n\t"
                 "mapa.shared::cluster.u32  remAddr32, %0, %1;\n\t"
                 "mbarrier.arrive.shared::cluster.b64  _, [remAddr32];\n\t"
                 "}"
                 :
                 : "r"(smem_int_ptr), "r"(cta_id));
  }
}

TL_DEVICE void mbarrier_expect_tx(uint64_t &smem_barrier,
                                  uint32_t transaction_bytes) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  asm volatile("mbarrier.expect_tx.shared.b64 [%1], %0;"
               :
               : "r"(transaction_bytes), "r"(smem_int_ptr));
}

TL_DEVICE void mbarrier_arrive_expect_tx(uint64_t &smem_barrier,
                                         uint32_t transaction_bytes) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%1], %0;"
               :
               : "r"(transaction_bytes), "r"(smem_int_ptr));
}

TL_DEVICE void mbarrier_cp_async_arrive(uint64_t &smem_barrier) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  asm volatile("cp.async.mbarrier.arrive.shared.b64 [%0];"
               :
               : "r"(smem_int_ptr));
}

TL_DEVICE void fence_proxy_async() {
  asm volatile("fence.proxy.async.shared::cta;" : :);
}

// Indicate arrival of warp issuing TMA_STORE
TL_DEVICE void tma_store_arrive() {
  asm volatile("cp.async.bulk.commit_group;");
}

template <int Count> TL_DEVICE void tma_store_wait() {
  asm volatile("cp.async.bulk.wait_group.read %0;" : : "n"(Count) : "memory");
}

TL_DEVICE void syncthreads_partial(uint64_t &smem_barrier) {
  uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
  uint64_t state = 0;
  asm volatile("{\n"
               ".reg .pred                P1;\n"
               "mbarrier.arrive.shared.b64 %1, [%0];\n"
               "LAB_WAIT:\n"
               "mbarrier.try_wait.shared.b64 P1, [%0], %1;\n"
               "@!P1                      bra.uni LAB_WAIT;\n"
               "}\n"
               :
               : "r"(smem_int_ptr), "l"(state));
}

template <uint32_t RegCount> TL_DEVICE void warpgroup_reg_alloc() {
  asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount));
}

template <uint32_t RegCount> TL_DEVICE void warpgroup_reg_dealloc() {
  asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount));
}

} // namespace tl