J3ff/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ 0798119

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:42 -05:00
commit b42569dbca
4529 changed files with 1132125 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/memory.cuh Normal file
View File

@@ -0,0 +1,10 @@
/**
* @file
* @brief An aggregate header of warp memory operations, where a single warp loads or stores data on its own.
*/
#pragma once
#include "util/util.cuh"
#include "tile/tile.cuh"
#include "vec/vec.cuh"

10
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/tile/tile.cuh Normal file
View File

@@ -0,0 +1,10 @@
/**
* @file
* @brief An aggregate header of warp memory operations on tiles, where a single warp loads or stores data on its own.
*/
#pragma once
#ifdef KITTENS_HOPPER
#include "tma.cuh"
#endif

564
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/tile/tma.cuh Normal file
View File

@@ -0,0 +1,564 @@
#pragma once
#include "../../../../common/common.cuh"
#include "../../../../types/types.cuh"
#include "../util/util.cuh"
#include <cuda.h>
#include <iostream>
namespace kittens {
namespace tma {
namespace detail {
template<kittens::ducks::st::all ST, int axis> __device__ inline int4 tma_coords(const coord<ducks::default_type> &unit_coord) {
constexpr int swizzle_elements = ST::swizzle_bytes / sizeof(typename ST::dtype);
if constexpr (axis == 2) return {unit_coord.r, unit_coord.c / swizzle_elements, unit_coord.d, unit_coord.b};
else if constexpr (axis == 1) return {unit_coord.d, unit_coord.c / swizzle_elements, unit_coord.r, unit_coord.b};
else if constexpr (axis == 0) return {unit_coord.b, unit_coord.c / swizzle_elements, unit_coord.r, unit_coord.d};
}
}
/* ---------- Prefetch Tensor Map ---------- */
/**
* @brief Prefetches data from global memory into a shared memory tile, along with the tensormap.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination shared memory tile.
* @param[in] src_tma_map The source tensormap address in global memory
* @param[in] tile_row_idx The row coord of the requested tile. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the requested tile. This is in units of complete tiles.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<ST, axis>());
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.prefetch.tensor.5d.L2.global.tile"
" [%0, {%1, %2, %3, %4, %5}];"
:
: "l"(tma_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.prefetch.tensor.5d.L2.global.tile.L2::cache_hint"
" [%0, {%1, %2, %3, %4, %5}], %6;"
:
: "l"(tma_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
prefetch<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
/* ---------- Async load and store data from gmem/smem ---------- */
/**
* @brief Asynchronously stores data into global memory from a shared memory tile.
*
* This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination tensormap address in global memory
* @param[in] src_tma_map The source shared memory tile.
* @param[in] tile_row_idx The row coord of the tile destination. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the tile destination. This is in units of complete tiles.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
store_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
store_async<dim::ROW, cache_policy::NORMAL>(dst, src, idx);
}
/* ---------- Async reduction + store data from gmem/smem ---------- */
/**
* @brief Asynchronously performs an add reduction and stores the result into global memory from a shared memory tile.
*
* This function performs an asynchronous add reduction and copy operation using CUDA's cp.reduce.async.bulk.tensor instruction.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination tensormap address in global memory
* @param[in] src_tma_map The source shared memory tile.
* @param[in] tile_row_idx The row coord of the tile destination. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the tile destination. This is in units of complete tiles.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
static_assert(!(std::is_same_v<typename ST::dtype, fp8e4m3> ||
std::is_same_v<typename ST::dtype, fp8e5m2>),
"TMA does not support async add reductions for fp8 types.");
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
store_add_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
static_assert(!(std::is_same_v<typename ST::dtype, fp8e4m3> ||
std::is_same_v<typename ST::dtype, fp8e5m2>),
"TMA does not support async add reductions for fp8 types.");
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
store_add_async<dim::ROW, cache_policy::NORMAL>(dst, src, idx);
}
/**
* @brief Asynchronously performs an min reduction and stores the result into global memory from a shared memory tile.
*
* This function performs an asynchronous min reduction and copy operation using CUDA's cp.reduce.async.bulk.tensor instruction.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination tensormap address in global memory
* @param[in] src_tma_map The source shared memory tile.
* @param[in] tile_row_idx The row coord of the tile destination. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the tile destination. This is in units of complete tiles.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
static_assert(!std::is_same_v<typename ST::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
static_assert(!(std::is_same_v<typename ST::dtype, fp8e4m3> ||
std::is_same_v<typename ST::dtype, fp8e5m2>),
"TMA does not support async add reductions for fp8 types.");
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
store_min_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
static_assert(!std::is_same_v<typename ST::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
static_assert(!(std::is_same_v<typename ST::dtype, fp8e4m3> ||
std::is_same_v<typename ST::dtype, fp8e5m2>),
"TMA does not support async add reductions for fp8 types.");
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.min.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
store_min_async<dim::ROW, cache_policy::NORMAL>(dst, src, idx);
}
/**
* @brief Asynchronously performs an max reduction and stores the result into global memory from a shared memory tile.
*
* This function performs an asynchronous max reduction and copy operation using CUDA's cp.reduce.async.bulk.tensor instruction.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination tensormap address in global memory
* @param[in] src_tma_map The source shared memory tile.
* @param[in] tile_row_idx The row coord of the tile destination. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the tile destination. This is in units of complete tiles.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
static_assert(!std::is_same_v<typename ST::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
static_assert(!(std::is_same_v<typename ST::dtype, fp8e4m3> ||
std::is_same_v<typename ST::dtype, fp8e5m2>),
"TMA does not support async add reductions for fp8 types.");
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
store_max_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
static_assert(!std::is_same_v<typename ST::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
static_assert(!(std::is_same_v<typename ST::dtype, fp8e4m3> ||
std::is_same_v<typename ST::dtype, fp8e5m2>),
"TMA does not support async add reductions for fp8 types.");
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<ST, axis>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.max.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(tma_ptr), "r"(src_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
store_commit_group();
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
store_max_async<dim::ROW, cache_policy::NORMAL>(dst, src, idx);
}
/**
* @brief Asynchronously loads data from global memory into a shared memory tile.
*
* This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination shared memory tile.
* @param[in] src_tma_map The source tensormap address in global memory
* @param[in,out] bar The semaphore used for synchronization of the asynchronous copy.
* @param[in] tile_row_idx The row coord of the requested tile. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the requested tile. This is in units of complete tiles.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<ST, axis>());
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
if constexpr (policy == cache_policy::NORMAL) {
asm volatile(
"cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2];"
:
: "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w)
: "memory"
);
}
else {
asm volatile(
"cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar);
}
namespace cluster {
/**
* @brief Asynchronously loads data from global memory into a shared memory tile, across a threadblock cluster
*
* This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
*
* @tparam ST A shared tile type with a TMA-compatible layout
* @param[out] dst The destination shared memory tile.
* @param[in] src_tma_map The source tensormap address in global memory
* @param[in,out] bar The semaphore used for synchronization of the asynchronous copy.
* @param[in] tile_row_idx The row coord of the requested tile. This is in units of complete tiles.
* @param[in] tile_col_idx The column coord of the requested tile. This is in units of complete tiles.
* @param[in] cluster_mask The mask of the clusters to broadcast to.
*/
#ifdef KITTENS_BLACKWELL
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1)
#else
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask)
#endif
{
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<ST, axis>());
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
coord<ducks::default_type> unit_coord = idx.template unit_coord<axis, 3>(); // convert to unit coordinates
int4 tma_coords = detail::tma_coords<ST, axis>(unit_coord);
#ifdef KITTENS_BLACKWELL
if(dst_mbar_cta != -1) {
uint32_t neighbor_mbar_ptr;
asm volatile (
"mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(neighbor_mbar_ptr)
: "r"(mbar_ptr), "r"(dst_mbar_cta)
);
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8, %9;"
:
: "r"(dst_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask), "l"(make_cache_policy<policy>())
: "memory"
);
}
} else
#endif
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8, %9;"
:
: "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr),
"n"(0), "r"(tma_coords.x), "r"(tma_coords.y), "r"(tma_coords.z), "r"(tma_coords.w), "h"(cluster_mask), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
#ifdef KITTENS_BLACKWELL
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
}
#else
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
}
#endif
} // namespace cluster
} // namespace tma
} // namespace kittens

405
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/util/multimem.cuh Normal file
View File

@@ -0,0 +1,405 @@
/**
* @file
* @brief Wrappers for multimem operations
*/
#pragma once
namespace kittens {
enum class reduce_op {
ADD = 0,
MIN = 1,
MAX = 2
};
enum class memory_model {
WEAK = 0,
STRONG = 1
};
template <typename T>
struct multimem;
template <>
struct multimem<int> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(int &dst, const int *src) {
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.s32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.s32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MIN) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.min.s32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.min.s32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MAX) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.max.s32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.max.s32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(int *dst, const int &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.s32 [%0], %1;"
:: "l"(dst), "r"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.s32 [%0], %1;"
:: "l"(dst), "r"(src) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(int *dst, const int &src) {
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.s32 [%0], %1;"
: : "l"(dst), "r"(src) : "memory");
} else if constexpr (Op == reduce_op::MIN) {
asm volatile("multimem.red.release.sys.global.min.s32 [%0], %1;"
: : "l"(dst), "r"(src) : "memory");
} else if constexpr (Op == reduce_op::MAX) {
asm volatile("multimem.red.release.sys.global.max.s32 [%0], %1;"
: : "l"(dst), "r"(src) : "memory");
}
}
};
template <>
struct multimem<uint> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(uint &dst, const uint *src) {
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.u32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.u32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MIN) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.min.u32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.min.u32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MAX) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.max.u32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.max.u32 %0, [%1];"
: "=r"(dst) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(uint *dst, const uint &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.u32 [%0], %1;"
:: "l"(dst), "r"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.u32 [%0], %1;"
:: "l"(dst), "r"(src) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(uint *dst, const uint &src) {
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.u32 [%0], %1;"
: : "l"(dst), "r"(src) : "memory");
} else if constexpr (Op == reduce_op::MIN) {
asm volatile("multimem.red.release.sys.global.min.u32 [%0], %1;"
: : "l"(dst), "r"(src) : "memory");
} else if constexpr (Op == reduce_op::MAX) {
asm volatile("multimem.red.release.sys.global.max.u32 [%0], %1;"
: : "l"(dst), "r"(src) : "memory");
}
}
};
template <>
struct multimem<float> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(float &dst, const float *src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 ld_reduce operations");
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.f32 %0, [%1];"
: "=f"(dst) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.f32 %0, [%1];"
: "=f"(dst) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(float *dst, const float &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.f32 [%0], %1;"
:: "l"(dst), "f"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.f32 [%0], %1;"
:: "l"(dst), "f"(src) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(float *dst, const float &src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 red operations");
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.f32 [%0], %1;"
: : "l"(dst), "f"(src) : "memory");
}
}
};
template <>
struct multimem<float2> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(float2 &dst, const float2 *src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 ld_reduce operations");
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.v2.f32 {%0, %1}, [%2];"
: "=f"(dst.x), "=f"(dst.y) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.v2.f32 {%0, %1}, [%2];"
: "=f"(dst.x), "=f"(dst.y) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(float2 *dst, const float2 &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.v2.f32 [%0], {%1, %2};"
:: "l"(dst), "f"(src.x), "f"(src.y) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.v2.f32 [%0], {%1, %2};"
:: "l"(dst), "f"(src.x), "f"(src.y) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(float2 *dst, const float2 &src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f32 red operations");
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.v2.f32 [%0], {%1, %2};"
: : "l"(dst), "f"(src.x), "f"(src.y) : "memory");
}
}
};
template <>
struct multimem<bf16> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(bf16 &dst, const bf16 *src) {
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.bf16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.bf16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MIN) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.min.bf16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.min.bf16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MAX) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.max.bf16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.max.bf16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(bf16 *dst, const bf16 &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.bf16 [%0], %1;"
:: "l"(dst), "h"(*reinterpret_cast<const uint16_t *>(&src)) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.bf16 [%0], %1;"
:: "l"(dst), "h"(*reinterpret_cast<const uint16_t *>(&src)) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(bf16 *dst, const bf16 &src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for bf16 red operations");
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.bf16 [%0], %1;"
: : "l"(dst), "h"(*reinterpret_cast<const uint16_t *>(&src)) : "memory");
}
}
};
template <>
struct multimem<bf16_2> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(bf16_2 &dst, const bf16_2 *src) {
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.bf16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.bf16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MIN) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.min.bf16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.min.bf16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MAX) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.max.bf16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.max.bf16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(bf16_2 *dst, const bf16_2 &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.bf16x2 [%0], %1;"
:: "l"(dst), "r"(*reinterpret_cast<const uint32_t *>(&src)) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.bf16x2 [%0], %1;"
:: "l"(dst), "r"(*reinterpret_cast<const uint32_t *>(&src)) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(bf16_2 *dst, const bf16_2 &src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for bf16_2 red operations");
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.bf16x2 [%0], %1;"
: : "l"(dst), "r"(*reinterpret_cast<const uint32_t *>(&src)) : "memory");
}
}
};
template <>
struct multimem<half> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(half &dst, const half *src) {
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.f16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.f16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MIN) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.min.f16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.min.f16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MAX) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.max.f16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.max.f16 %0, [%1];"
: "=h"(*reinterpret_cast<uint16_t *>(&dst)) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(half *dst, const half &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.f16 [%0], %1;"
:: "l"(dst), "h"(*reinterpret_cast<const uint16_t *>(&src)) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.f16 [%0], %1;"
:: "l"(dst), "h"(*reinterpret_cast<const uint16_t *>(&src)) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(half *dst, const half &src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f16 red operations");
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.f16 [%0], %1;"
: : "l"(dst), "h"(*reinterpret_cast<const uint16_t *>(&src)) : "memory");
}
}
};
template <>
struct multimem<half_2> {
template <reduce_op Op, memory_model M = memory_model::WEAK>
__device__ static inline void ld_reduce(half_2 &dst, const half_2 *src) {
if constexpr (Op == reduce_op::ADD) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.f16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.add.acc::f32.f16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MIN) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.min.f16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.min.f16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
}
} else if constexpr (Op == reduce_op::MAX) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.ld_reduce.weak.global.max.f16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.ld_reduce.acquire.sys.global.max.f16x2 %0, [%1];"
: "=r"(*reinterpret_cast<uint32_t *>(&dst)) : "l"(src) : "memory");
}
}
}
template <memory_model M = memory_model::WEAK>
__device__ static inline void st(half_2 *dst, const half_2 &src) {
if constexpr (M == memory_model::WEAK) {
asm volatile("multimem.st.weak.global.f16x2 [%0], %1;"
:: "l"(dst), "r"(*reinterpret_cast<const uint32_t *>(&src)) : "memory");
} else if constexpr (M == memory_model::STRONG) {
asm volatile("multimem.st.release.sys.global.f16x2 [%0], %1;"
:: "l"(dst), "r"(*reinterpret_cast<const uint32_t *>(&src)) : "memory");
}
}
template <reduce_op Op>
__device__ static inline void red(half_2 *dst, const half_2 &src) {
static_assert(Op == reduce_op::ADD, "MIN/MAX are not supported for f16_2 red operations");
if constexpr (Op == reduce_op::ADD) {
asm volatile("multimem.red.release.sys.global.add.f16x2 [%0], %1;"
: : "l"(dst), "r"(*reinterpret_cast<const uint32_t *>(&src)) : "memory");
}
}
};
} // namespace kittens

30
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/util/tensor.cuh Normal file
View File

@@ -0,0 +1,30 @@
/**
* @file
* @brief Functions for transferring data directly between tensor memory and register memory.
*/
#pragma once
#include <type_traits>
#include "../../../../common/common.cuh"
#include "../../../../types/types.cuh"
#include "util.cuh"
namespace kittens {
__device__ static inline void tensor_before_thread_sync() {
asm volatile("tcgen05.fence::before_thread_sync;\n");
}
__device__ static inline void tensor_after_thread_sync() {
asm volatile("tcgen05.fence::after_thread_sync;\n");
}
__device__ inline static void tensor_load_wait() {
asm volatile("tcgen05.wait::ld.sync.aligned;");
}
__device__ inline static void tensor_store_wait() {
asm volatile("tcgen05.wait::st.sync.aligned;");
}
}

249
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/util/tma.cuh Normal file
View File

@@ -0,0 +1,249 @@
#pragma once
#include "../../../../common/common.cuh"
#include "../../../../types/types.cuh"
#include <cuda.h>
#include <iostream>
namespace kittens {
/**
* @brief A namespace for all of ThunderKittens' TMA functionality.
*/
namespace tma {
/* ---------- Barrier functions for async load ---------- */
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
* instruction to set the expected number of bytes.
*
* @param semaphore Reference to the semaphore variable.
* @param bytes The number of bytes expected at the semaphore.
*/
__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) {
void const* const ptr = &bar;
uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n"
:: "r"(bar_ptr), "r"(bytes));
}
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
*/
template<typename T, typename... args>
__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) {
expect_bytes(bar, size_bytes<T, args...>);
}
/* ---------- Synchronization functions for async store ---------- */
/**
* @brief Commits previous asynchronous TMA stores to a group and performs them.
*/
__device__ static inline void store_commit_group() {
asm volatile("cp.async.bulk.commit_group;");
}
/**
* @brief Waits for previous committed TMA store groups to complete.
*
* @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
*/
template <int N=0>
__device__ static inline void store_async_wait() {
asm volatile (
"cp.async.bulk.wait_group %0;"
:
: "n"(N)
: "memory"
);
}
/**
* @brief Waits for previous committed TMA store groups to finish reading from shared memory.
*
* @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
*/
template <int N=0>
__device__ static inline void store_async_read_wait() {
asm volatile (
"cp.async.bulk.wait_group.read %0;"
:
: "n"(N)
: "memory"
);
}
/* ---------- Cluster-scope operations ---------- */
namespace cluster {
/**
* @brief Waits for the requested semaphore phase, at cluster scope
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void wait(semaphore& bar, int kPhaseBit) {
void const* const ptr = &bar;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
}
__device__ static inline void careful_wait(semaphore& bar, int kPhaseBit) {
void const* const ptr = &bar;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"{\n"
".reg .b64 start_clock, current_clock;\n"
"mov.b64 start_clock, %clock64;\n"
".reg .pred P_CLOCK;\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"mov.b64 current_clock, %clock64;\n"
"sub.u64 current_clock, current_clock, start_clock;\n"
"setp.ge.u64 P_CLOCK, current_clock, 1000000;\n"
"@P_CLOCK trap;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
}
/**
* @brief Sets the number of bytes expected at the semaphore, assuming a multicast instruction.
*
* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
* instruction to set the expected number of bytes.
*
* It's worth being aware that this function is particularly necessary for multicast loads, and
* distributed shared memory can actually be done with a normal tma::expect followed by wait. See
* the unit tests of dsmem for an example.
*
* @param semaphore Reference to the semaphore variable.
* @param bytes The number of bytes expected at the semaphore.
*/
__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) {
uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t neighbor_mbar_addr;
asm volatile (
"mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(neighbor_mbar_addr)
: "r"(mbar_addr), "r"(dst_cta)
);
asm volatile ("mbarrier.arrive.expect_tx.shared::cluster.b64 _, [%0], %1;\n"
:: "r"(neighbor_mbar_addr), "r"(bytes));
}
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
* instruction to set the expected number of bytes.
*
* @tparam T The type of the data to be stored at the semaphore.
* @param semaphore Reference to the semaphore variable.
*/
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
*/
template<typename T, typename... args>
__device__ static inline void expect(semaphore& bar, int dst_cta, const T& _1, const args&... _2) {
expect_bytes(bar, size_bytes<T, args...>, dst_cta);
}
/**
* @brief Arrives at a semaphore in cluster scope.
*
* Marks a thread arrival at an mbarrier
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) {
uint32_t mbar_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t neighbor_mbar_addr;
asm volatile (
"mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(neighbor_mbar_addr)
: "r"(mbar_addr), "r"(dst_cta)
);
asm volatile (
"mbarrier.arrive.shared::cluster.b64 _, [%0], %1;\n"
:
: "r"(neighbor_mbar_addr), "r" (count)
: "memory"
);
}
// Generic transfer
__device__ static inline void store_async(void *dst, void *src, int dst_cta, uint32_t size_bytes, semaphore& bar) {
void const* const ptr = &bar;
uint32_t mbarrier_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
// **************************************************
// load from src to dst in different threadblocks
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(src));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(dst));
// mapa instr = https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-mapa
// find dst addr in neighbor's cta
uint32_t neighbor_addr_dst;
asm volatile (
"mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(neighbor_addr_dst)
: "r"(dst_ptr), "r"(dst_cta)
);
uint32_t neighbor_addr_mbarrier = mbarrier_ptr;
asm volatile (
"mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(neighbor_addr_mbarrier)
: "r"(mbarrier_ptr), "r"(dst_cta)
);
// cp.async instr = https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk
// copy src into dst in neighbor's cta
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
asm volatile (
"cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n"
:
: "r"(neighbor_addr_dst), "r"(src_ptr), "r"(size_bytes), "r"(neighbor_addr_mbarrier)
: "memory"
);
}
// Templated transfer for convenience
template<typename T>
__device__ static inline void store_async(T &dst_, T &src_, int dst_cta, semaphore& bar) {
store_async((void*)&dst_, (void*)&src_, dst_cta, size_bytes<T>, bar);
}
} // namespace cluster
} // namespace tma
} // namespace kittens

443
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/util/util.cuh Normal file
View File

@@ -0,0 +1,443 @@
/**
* @file
* @brief General memory utilities not specialized for either tiles or vectors.
*/
#pragma once
namespace kittens {
/* ---------- To prevent generic addressing, PTX ---------- */
template<typename T> struct move {
__device__ static inline void lds(T& dst, uint32_t src);
__device__ static inline void sts(uint32_t dst, const T& src);
__device__ static inline void ldg(T& dst, T* src);
__device__ static inline void stg(T* dst, const T& src);
};
// unpacked types
template<> struct move<bf16> {
__device__ static inline void lds(bf16& dst, uint32_t src) {
asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const bf16& src) {
asm volatile("st.shared.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "r"(dst));
}
__device__ static inline void ldg(bf16& dst, bf16* src) {
asm volatile("ld.global.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "l"(src));
}
__device__ static inline void stg(bf16* dst, const bf16& src) {
asm volatile("st.global.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "l"(dst));
}
};
template<> struct move<half> {
__device__ static inline void lds(half& dst, uint32_t src) {
asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const half& src) {
asm volatile("st.shared.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "r"(dst));
}
__device__ static inline void ldg(half& dst, half* src) {
asm volatile("ld.global.b16 %0, [%1];\n" : "=h"(*(uint16_t*)&dst) : "l"(src));
}
__device__ static inline void stg(half* dst, const half& src) {
asm volatile("st.global.b16 [%1], %0;\n" : : "h"(*(uint16_t*)&src), "l"(dst));
}
};
template<> struct move<float> {
__device__ static inline void lds(float& dst, uint32_t src) {
asm volatile("ld.shared.f32 %0, [%1];\n" : "=f"(dst) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const float& src) {
asm volatile("st.shared.f32 [%1], %0;\n" : : "f"(src), "r"(dst));
}
__device__ static inline void ldg(float& dst, float* src) {
asm volatile("ld.global.f32 %0, [%1];\n" : "=f"(dst) : "l"(src));
}
__device__ static inline void stg(float* dst, const float& src) {
asm volatile("st.global.f32 [%1], %0;\n" : : "f"(src), "l"(dst));
}
};
template<> struct move<int> {
__device__ static inline void lds(int& dst, uint32_t src) {
asm volatile("ld.shared.u32 %0, [%1];\n" : "=r"(dst) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const int& src) {
asm volatile("st.shared.u32 [%1], %0;\n" : : "r"(src), "r"(dst));
}
__device__ static inline void ldg(int& dst, int* src) {
asm volatile("ld.global.u32 %0, [%1];\n" : "=r"(dst) : "l"(src));
}
__device__ static inline void stg(int* dst, const int& src) {
asm volatile("st.global.u32 [%1], %0;\n" : : "r"(src), "l"(dst));
}
};
// packed types
template<> struct move<bf16_2> {
__device__ static inline void lds(bf16_2& dst, uint32_t src) {
asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const bf16_2& src) {
asm volatile("st.shared.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "r"(dst));
}
__device__ static inline void ldg(bf16_2& dst, bf16_2* src) {
asm volatile("ld.global.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "l"(src));
}
__device__ static inline void stg(bf16_2* dst, const bf16_2& src) {
asm volatile("st.global.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "l"(dst));
}
__device__ static inline void ldsm4(bf16_2& dst1, bf16_2& dst2, bf16_2& dst3, bf16_2& dst4, uint32_t src) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" :
"=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src));
}
__device__ static inline void ldsm4t(bf16_2& dst1, bf16_2& dst2, bf16_2& dst3, bf16_2& dst4, uint32_t src) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" :
"=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src));
}
__device__ static inline void stsm4(uint32_t dst, bf16_2& src1, bf16_2& src2, bf16_2& src3, bf16_2& src4) {
asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" ::
"r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst));
}
__device__ static inline void stsm4t(uint32_t dst, bf16_2& src1, bf16_2& src2, bf16_2& src3, bf16_2& src4) {
asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" ::
"r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst));
}
};
template<> struct move<half_2> {
__device__ static inline void lds(half_2& dst, uint32_t src) {
asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const half_2& src) {
asm volatile("st.shared.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "r"(dst));
}
__device__ static inline void ldg(half_2& dst, half_2* src) {
asm volatile("ld.global.b32 %0, [%1];\n" : "=r"(*(uint32_t*)&dst) : "l"(src));
}
__device__ static inline void stg(half_2* dst, const half_2& src) {
asm volatile("st.global.b32 [%1], %0;\n" : : "r"(*(uint32_t*)&src), "l"(dst));
}
__device__ static inline void ldsm4(half_2& dst1, half_2& dst2, half_2& dst3, half_2& dst4, uint32_t src) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" :
"=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src));
}
__device__ static inline void ldsm4t(half_2& dst1, half_2& dst2, half_2& dst3, half_2& dst4, uint32_t src) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" :
"=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src));
}
__device__ static inline void stsm4(uint32_t dst, half_2& src1, half_2& src2, half_2& src3, half_2& src4) {
asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" ::
"r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst));
}
__device__ static inline void stsm4t(uint32_t dst, half_2& src1, half_2& src2, half_2& src3, half_2& src4) {
asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" ::
"r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst));
}
};
template<> struct move<float2> {
__device__ static inline void lds(float2& dst, uint32_t src) {
asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];\n" : "=f"(dst.x), "=f"(dst.y) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const float2& src) {
asm volatile("st.shared.v2.f32 [%2], {%0, %1};\n" : : "f"(src.x), "f"(src.y), "r"(dst));
}
__device__ static inline void ldg(float2& dst, float2* src) {
asm volatile("ld.global.v2.f32 {%0, %1}, [%2];\n" : "=f"(dst.x), "=f"(dst.y) : "l"(src));
}
__device__ static inline void stg(float2* dst, const float2& src) {
asm volatile("st.global.v2.f32 [%2], {%0, %1};\n" : : "f"(src.x), "f"(src.y), "l"(dst));
}
};
template<> struct move<float4> {
__device__ static inline void lds(float4& dst, uint32_t src) {
asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];\n" : "=f"(dst.x), "=f"(dst.y), "=f"(dst.z), "=f"(dst.w) : "r"(src));
}
__device__ static inline void sts(uint32_t dst, const float4& src) {
asm volatile("st.shared.v4.f32 [%4], {%0, %1, %2, %3};\n" : : "f"(src.x), "f"(src.y), "f"(src.z), "f"(src.w), "r"(dst));
}
__device__ static inline void ldg(float4& dst, float4* src) {
asm volatile("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" : "=f"(dst.x), "=f"(dst.y), "=f"(dst.z), "=f"(dst.w) : "l"(src));
}
__device__ static inline void stg(float4* dst, const float4& src) {
asm volatile("st.global.v4.f32 [%4], {%0, %1, %2, %3};\n" : : "f"(src.x), "f"(src.y), "f"(src.z), "f"(src.w), "l"(dst));
}
};
#ifdef KITTENS_HOPPER
template<> struct move<fp8e4m3_4> {
__device__ static inline void ldsm4(fp8e4m3_4& dst1, fp8e4m3_4& dst2, fp8e4m3_4& dst3, fp8e4m3_4& dst4, uint32_t src) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" :
"=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src));
}
__device__ static inline void stsm4(uint32_t dst, fp8e4m3_4& src1, fp8e4m3_4& src2, fp8e4m3_4& src3, fp8e4m3_4& src4) {
asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" ::
"r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst));
}
};
template<> struct move<fp8e5m2_4> {
__device__ static inline void ldsm4(fp8e5m2_4& dst1, fp8e5m2_4& dst2, fp8e5m2_4& dst3, fp8e5m2_4& dst4, uint32_t src) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n" :
"=r"(*(uint32_t*)&dst1), "=r"(*(uint32_t*)&dst2), "=r"(*(uint32_t*)&dst3), "=r"(*(uint32_t*)&dst4) : "r"(src));
}
__device__ static inline void stsm4(uint32_t dst, fp8e5m2_4& src1, fp8e5m2_4& src2, fp8e5m2_4& src3, fp8e5m2_4& src4) {
asm volatile("stmatrix.sync.aligned.m8n8.x4.shared::cta.b16 [%4], {%0, %1, %2, %3};\n" ::
"r"(*(uint32_t*)&src1), "r"(*(uint32_t*)&src2), "r"(*(uint32_t*)&src3), "r"(*(uint32_t*)&src4), "r"(dst));
}
};
#endif
/* ---------- Constants for Cache policies ---------- */
enum cache_policy {
NORMAL = 0,
EVICT_FIRST = 1,
EVICT_LAST = 2
};
template<cache_policy policy> __device__ inline uint64_t make_cache_policy() {
uint64_t cache_policy_val;
constexpr float fraction = 1.0f;
static_assert(policy == cache_policy::EVICT_FIRST || policy == cache_policy::EVICT_LAST, "Unexpected cache policy");
if constexpr (policy == cache_policy::EVICT_FIRST) {
asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, %1;\n" : "=l"(cache_policy_val) : "f"(fraction));
}
else {
asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, %1;\n" : "=l"(cache_policy_val) : "f"(fraction));
}
return cache_policy_val;
}
/* ---------- Generic (non-Hopper specific) semaphore functions ---------- */
struct semaphore {
private:
uint64_t value;
}; // note that this is an opaque type, so the value should not be accessed directly.
template<int num_warps> struct barrier {
int barrier_id;
__device__ __forceinline__ barrier(int _id) : barrier_id(_id) {}
__device__ __forceinline__ barrier operator[](int i) {
return barrier(barrier_id + i);
}
};
/**
* @brief Initializes a synchronization semaphore with a transaction count and sets the expected number of bytes.
*
* This function sets up a semaphore that is used to synchronize threads within a block during asynchronous operations.
* It initializes the semaphore with a thread count semaphore.
*
* Additionally, if it is given a shared tile type, it will also call `set_bytes` to prepare for the memory transaction.
*
* @param[out] semaphore The semaphore variable to initialize.
* @param[in] tc The thread counter for the semaphore.
*/
__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) {
void const* const ptr = &bar;
uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"mbarrier.init.shared::cta.b64 [%0], %1;\n"
:: "r"(bar_ptr), "r"(thread_count+transaction_count)
);
}
/**
* @brief Invalidate an mbarrier
*
* @param[out] semaphore The semaphore variable to initialize.
* @param[in] tc The thread counter for the semaphore.
*/
__device__ static inline void invalidate_semaphore(semaphore& bar) {
void const* const ptr = &bar;
uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"mbarrier.inval.shared::cta.b64 [%0];\n"
:: "r"(bar_ptr)
);
}
/**
* @brief Arrives at a semaphore.
*
* Marks a warp arrival at an mbarrier
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void arrive(semaphore& sem) {
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem));
asm volatile (
"mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n"
:
: "r"(mbar_ptr)
: "memory"
);
}
template<int num_warps> __device__ static inline void arrive(barrier<num_warps> bar) {
asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory");
}
#ifdef KITTENS_HOPPER
/**
* @brief Arrives at a semaphore.
*
* Marks a warp arrival at an mbarrier
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void arrive(semaphore& sem, uint32_t count) {
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem));
asm volatile (
"mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n"
:
: "r"(mbar_ptr), "r"(count)
: "memory"
);
}
#endif
/**
* @brief Waits for the requested semaphore phase.
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void wait(semaphore& sem, int kPhaseBit) {
void const* const ptr = &sem;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
#ifdef KITTENS_HOPPER
asm volatile (
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
#else
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"(mbar_ptr),
"r"(kPhaseBit)
);
#endif
}
__device__ static inline void careful_wait(semaphore& sem, int kPhaseBit) {
void const* const ptr = &sem;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
#ifdef KITTENS_HOPPER
asm volatile (
"{\n"
".reg .b64 start_clock, current_clock;\n"
"mov.b64 start_clock, %clock64;\n"
".reg .pred P_CLOCK;\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"mov.b64 current_clock, %clock64;\n"
"sub.u64 current_clock, current_clock, start_clock;\n"
"setp.ge.u64 P_CLOCK, current_clock, 1000000;\n"
"@P_CLOCK trap;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
#else
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"(mbar_ptr),
"r"(kPhaseBit)
);
#endif
}
/**
* @brief Checks if the requested semaphore phase is ready.
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) {
void const* const ptr = &sem;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
int result;
asm volatile (
"{\n"
".reg .pred P1;\n"
"mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n"
"selp.u32 %0,1,0,P1;"
"}\n"
: "=r"(result)
: "r"(mbar_ptr), "r"(kPhaseBit)
);
return result;
}
__device__ static inline void arrive_and_wait(semaphore& sem, int kPhaseBit) {
arrive(sem);
wait(sem, kPhaseBit);
}
template<int num_warps> __device__ static inline void arrive_and_wait(barrier<num_warps> bar) {
asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory");
}
template<int N=0> __device__ static inline void load_async_wait() { // for completing (non-TMA) async loads
if constexpr (N == 0) {
asm volatile("cp.async.wait_all;\n" ::);
} else {
asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}
__syncwarp();
}
// meant to be used only with shared tiles and shared vectors
namespace detail {
template<typename T> struct size_info {
static constexpr uint32_t bytes = sizeof(std::remove_reference_t<T>);
};
template<ducks::st::all ST> struct size_info<ST> {
static constexpr uint32_t elements = ST::num_elements;
static constexpr uint32_t bytes = ST::num_elements * sizeof(typename ST::dtype);
};
template<ducks::sv::all SV> struct size_info<SV> {
static constexpr uint32_t elements = SV::length;
static constexpr uint32_t bytes = SV::length * sizeof(typename SV::dtype);
};
}
template<typename... Args> inline constexpr uint32_t size_bytes = 0; // base case
template<typename T, typename... Args> inline constexpr uint32_t size_bytes<T, Args...> = detail::size_info<T>::bytes + size_bytes<Args...>; // recursive case
} // namespace kittens
#ifdef KITTENS_HOPPER
#include "multimem.cuh"
#include "tma.cuh"
#endif
#ifdef KITTENS_BLACKWELL
#include "tensor.cuh"
#endif

416
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/vec/tma.cuh Normal file
View File

@@ -0,0 +1,416 @@
#pragma once
#include "../../../../common/common.cuh"
#include "../../../../types/types.cuh"
#include "../util/util.cuh"
#include <cuda.h>
#include <iostream>
// This is a macro that helps us define default cache policy versions of each function.
#define __KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx) { \
function_name<cache_policy::NORMAL>(dst, src, idx); \
}
#define __KITTENS_TMA_DEFINE_PGL_DEFAULT_LOAD_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx) { \
function_name<cache_policy::NORMAL>(dst, src, idx); \
}
#define __KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(const GL &dst, const SV &src, const COORD &idx) { \
function_name<cache_policy::NORMAL>(dst, src, idx); \
}
#define __KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(const PGL &dst, const SV &src, const COORD &idx) { \
function_name<cache_policy::NORMAL>(dst, src, idx); \
}
#define __KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx, semaphore& bar) { \
function_name<cache_policy::NORMAL>(dst, src, idx, bar); \
}
#define __KITTENS_TMA_DEFINE_PGL_SEMAPHORE_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx, semaphore& bar) { \
function_name<cache_policy::NORMAL>(dst, src, idx, bar); \
}
#define __KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { \
function_name<cache_policy::NORMAL>(dst, src, idx, bar, cluster_mask, dst_mbar_cta); \
}
#define __KITTENS_TMA_DEFINE_PGL_CLUSTER_SEMAPHORE_CACHE_VEC__(function_name) \
template<ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>> \
__device__ static inline void function_name(SV &dst, const PGL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) { \
function_name<cache_policy::NORMAL>(dst, src, idx, bar, cluster_mask, dst_mbar_cta); \
}
namespace kittens {
namespace detail {
namespace tma {
template<cache_policy policy> __device__ static inline void vec_prefetch_tma_internal(uint64_t tma_ptr, coord<> tma_coord) {
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.prefetch.tensor.4d.L2.global.tile"
" [%0, {%1, %2, %3, %4}];"
:
: "l"(tma_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.prefetch.tensor.4d.L2.global.tile.L2::cache_hint"
" [%0, {%1, %2, %3, %4}], %5;"
:
: "l"(tma_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<cache_policy policy> __device__ static inline void vec_store_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) {
asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group"
" [%0, {%2, %3, %4, %5}], [%1];"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5}], [%1], %6;"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<cache_policy policy> __device__ static inline void vec_store_add_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) {
asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group"
" [%0, {%2, %3, %4, %5}], [%1];"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5}], [%1], %6;"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<cache_policy policy> __device__ static inline void vec_store_min_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) {
asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group"
" [%0, {%2, %3, %4, %5}], [%1];"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.min.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5}], [%1], %6;"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<cache_policy policy> __device__ static inline void vec_store_max_async_tma_internal(uint64_t tma_ptr, uint32_t src_i_ptr, coord<> tma_coord) {
asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group"
" [%0, {%2, %3, %4, %5}], [%1];"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b)
: "memory"
);
}
else {
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.max.tile.bulk_group.L2::cache_hint"
" [%0, {%2, %3, %4, %5}], [%1], %6;"
:
: "l"(tma_ptr), "r"(src_i_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
template<cache_policy policy> __device__ static inline void vec_load_async_tma_internal(uint64_t tma_ptr, uint32_t dst_i_ptr, uint32_t mbar_ptr, coord<> tma_coord) {
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2];"
:
: "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
:
: "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr), "r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
namespace cluster {
template<cache_policy policy> __device__ static inline void vec_load_async_tma_internal(uint64_t tma_ptr, uint32_t dst_i_ptr, uint32_t mbar_ptr, coord<> tma_coord, uint16_t cluster_mask, int dst_mbar_cta=-1) {
#ifdef KITTENS_BLACKWELL
if(dst_mbar_cta != -1) {
uint32_t neighbor_mbar_ptr;
asm volatile (
"mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(neighbor_mbar_ptr)
: "r"(mbar_ptr), "r"(dst_mbar_cta)
);
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
:
: "r"(dst_i_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr),
"r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.cta_group::2.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;"
:
: "r"(dst_i_ptr), "l"(tma_ptr), "r"(neighbor_mbar_ptr),
"r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask), "l"(make_cache_policy<policy>())
: "memory"
);
}
} else
#endif
if constexpr (policy == cache_policy::NORMAL) {
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
:
: "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr),
"r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask)
: "memory"
);
}
else {
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7, %8;"
:
: "r"(dst_i_ptr), "l"(tma_ptr), "r"(mbar_ptr),
"r"(tma_coord.c), "r"(tma_coord.r), "r"(tma_coord.d), "r"(tma_coord.b), "h"(cluster_mask), "l"(make_cache_policy<policy>())
: "memory"
);
}
}
} // namespace cluster
} // namespace tma
} // namespace detail
namespace tma {
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
::kittens::detail::tma::vec_prefetch_tma_internal<policy>(tma_ptr, tma_coord);
}
}
__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch)
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async)
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_add_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_add_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async)
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_min_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_min_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async)
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_max_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_max_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
::kittens::tma::store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async)
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_load_async_tma_internal<policy>(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord);
}
}
__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async)
namespace cluster {
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
for(int i = 0; i < ::kittens::detail::tma::sv_tma_dim2<SV>; i++) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::cluster::vec_load_async_tma_internal<policy>(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord, cluster_mask, dst_mbar_cta);
}
}
__KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(load_async)
} // namespace cluster
} // namespace tma
} // namespace kittens

10
tinygrad_repo/extra/thunder/cuda/include/ops/thread/memory/vec/vec.cuh Normal file
View File

@@ -0,0 +1,10 @@
/**
* @file
* @brief An aggregate header of warp memory operations on vectors, where a single warp loads or stores data on its own.
*/
#pragma once
#ifdef KITTENS_HOPPER
#include "tma.cuh"
#endif

8
tinygrad_repo/extra/thunder/cuda/include/ops/thread/mma/mma.cuh Normal file
View File

@@ -0,0 +1,8 @@
/**
* @file
* @brief An aggregate header for warp operations on data stored in tensor memory.
*/
#pragma once
#include "tensor/tensor.cuh"

523
tinygrad_repo/extra/thunder/cuda/include/ops/thread/mma/tensor/tensor.cuh Normal file
View File

@@ -0,0 +1,523 @@
/**
* @file
* @brief Matrix multiply-accumulate operations for tiles stored in tensor memory.
*/
#pragma once
#include "../../../../common/common.cuh"
#include "../../../../types/types.cuh"
namespace kittens {
namespace detail {
namespace tcgen05 {
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#instruction-descriptor
template<typename D, typename AB, int M, int N, bool trans_a, bool trans_b, bool neg=false>
__device__ static inline uint32_t instruction_descriptor() {
uint32_t desc = 0;
if constexpr (sizeof(AB) == 2) { // kind::f16
// either accumulate to float, or the input is half and the output is half
static_assert(std::is_same_v<D, float> || std::is_same_v<AB, half>);
desc |= 0b00 << 0; // sparsity bits unneeded
desc |= 0b0 << 2; // dense
desc |= 0b0 << 3; // no saturate on fp types
if constexpr (std::is_same_v<D, float>) {
desc |= 0b01 << 4; // D matrix is FP32
}
else {
desc |= 0b00 << 4; // D matrix is FP16
}
desc |= 0b0 << 6; // reserved
if constexpr (std::is_same_v<AB, half>) {
desc |= 0b000 << 7; // 16-bit A input type as FP16
desc |= 0b000 << 10; // 16-bit B input type as FP16
} else if constexpr (std::is_same_v<AB, bf16>) {
desc |= 0b001 << 7; // 16-bit A input type as BF16
desc |= 0b001 << 10; // 16-bit B input type as BF16
} else if constexpr (std::is_same_v<AB, fp8e4m3>) {
desc |= 0b000 << 7; // 8-bit A input type as FP8 e4m3
desc |= 0b000 << 10; // 8-bit B input type as FP8 e4m3
} else if constexpr (std::is_same_v<AB, fp8e5m2>) {
desc |= 0b001 << 7; // 8-bit A input type as FP8 e5m2
desc |= 0b001 << 10; // 8-bit B input type as FP8 e5m2
}
/* fp6 and fp4
else if constexpr (std::is_same_v<AB, fp6e2m3>) {
desc |= 0b011 << 7; // 6-bit A input type as FP6 e2m3
desc |= 0b011 << 10; // 6-bit B input type as FP6 e2m3
}
else if constexpr (std::is_same_v<AB, fp4e2m3>) {
desc |= 0b100 << 7; // 6-bit A input type as FP6 e3m2
desc |= 0b100 << 10; // 6-bit B input type as FP6 e3m2
}
else if constexpr (std::is_same_v<AB, fp4e3m1>) {
desc |= 0b101 << 7; // 4-bit A input type as FP4 e3m1
desc |= 0b101 << 10; // 4-bit B input type as FP4 e3m1
}
*/
if constexpr (neg) {
desc |= 0b1 << 13; // Do negate A matrix
}
else {
desc |= 0b0 << 13; // Don't negate A matrix
}
desc |= 0b0 << 14; // Don't negate B matrix (in all cases)
if constexpr (trans_a) {
desc |= 0b1 << 15; // Transpose A matrix
}
else {
desc |= 0b0 << 15; // Don't transpose A matrix
}
if constexpr (trans_b) {
desc |= 0b1 << 16; // Transpose B matrix
}
else {
desc |= 0b0 << 16; // Don't transpose B matrix
}
desc |= (N >> 3) << 17; // B matrix has dimension N, encoded
desc |= 0b0 << 23; // reserved
desc |= (M >> 4) << 24; // A matrix has dimension M, encoded
desc |= 0b0 << 29; // reserved
desc |= 0b00 << 30; // no shift for B-matrix reuse
} else if constexpr (sizeof(AB) == 1) { // kind::f8f6f4
static_assert(std::is_same_v<D, float> || std::is_same_v<D, half>); // FP8/6/4 has to accumulate to float or half
desc |= 0b00 << 0; // sparsity bits unneeded
desc |= 0b0 << 2; // dense
desc |= 0b0 << 3; // no saturate on fp types
if constexpr (std::is_same_v<D, float>) {
desc |= 0b01 << 4; // D matrix is FP32
}
else {
desc |= 0b00 << 4; // D matrix is FP16
}
desc |= 0b0 << 6; // reserved
if constexpr (std::is_same_v<AB, fp8e4m3>) {
desc |= 0b000 << 7; // 8-bit A input type as FP8 e4m3
desc |= 0b000 << 10; // 8-bit B input type as FP8 e4m3
} else if constexpr (std::is_same_v<AB, fp8e5m2>) {
desc |= 0b001 << 7; // 8-bit A input type as FP8 e5m2
desc |= 0b001 << 10; // 8-bit B input type as FP8 e5m2
}
/* fp6 and fp4
else if constexpr (std::is_same_v<AB, fp6e2m3>) {
desc |= 0b011 << 7; // 6-bit A input type as FP6 e2m3
desc |= 0b011 << 10; // 6-bit B input type as FP6 e2m3
}
else if constexpr (std::is_same_v<AB, fp4e2m3>) {
desc |= 0b100 << 7; // 6-bit A input type as FP6 e3m2
desc |= 0b100 << 10; // 6-bit B input type as FP6 e3m2
}
else if constexpr (std::is_same_v<AB, fp4e3m1>) {
desc |= 0b101 << 7; // 4-bit A input type as FP4 e3m1
desc |= 0b101 << 10; // 4-bit B input type as FP4 e3m1
}
*/
if constexpr (neg) {
desc |= 0b1 << 13; // Do negate A matrix
}
else {
desc |= 0b0 << 13; // Don't negate A matrix
}
desc |= 0b0 << 14; // Don't negate B matrix (in all cases)
if constexpr (trans_a) {
desc |= 0b1 << 15; // Transpose A matrix
}
else {
desc |= 0b0 << 15; // Don't transpose A matrix
}
if constexpr (trans_b) {
desc |= 0b1 << 16; // Transpose B matrix
}
else {
desc |= 0b0 << 16; // Don't transpose B matrix
}
desc |= (N >> 3) << 17; // B matrix has dimension N, encoded
desc |= 0b0 << 23; // reserved
desc |= (M >> 4) << 24; // A matrix has dimension M, encoded
desc |= 0b0 << 29; // reserved
desc |= 0b00 << 30; // no shift for B-matrix reuse
}
else {
static_assert(sizeof(AB) == 999, "Invalid AB type size; not implemented yet.");
}
return desc;
};
template<typename T_AB, int acc, int ncta=1>
__device__ static inline void tt_st(uint32_t d_tt_addr, uint32_t a_tt_addr, uint64_t b_desc, uint32_t idesc) {
if constexpr (std::is_same_v<T_AB, fp8e4m3> || std::is_same_v<T_AB, fp8e5m2>) {
// TODO(danfu): is there a better way to do this with string manipulation that the compiler likes?
if constexpr (ncta == 1) {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [%1], %2, %3, p;}\n"
:: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
else {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], [%1], %2, %3, p;}\n"
:: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
} else {
if constexpr (ncta == 1) {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, p;}\n"
:: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
else {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::2.kind::f16 [%0], [%1], %2, %3, p;}\n"
:: "r"(d_tt_addr), "r"(a_tt_addr), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
}
}
template<typename T_AB, int acc, int ncta=1>
__device__ static inline void st_st(uint32_t d_tt_addr, uint64_t a_desc, uint64_t b_desc, uint32_t idesc) {
if constexpr (std::is_same_v<T_AB, fp8e4m3> || std::is_same_v<T_AB, fp8e5m2>) {
// TODO(danfu): is there a better way to do this with string manipulation that the compiler likes?
if constexpr (ncta == 1) {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;}\n"
:: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
else {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, p;}\n"
:: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
} else {
if constexpr (ncta == 1) {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;}\n"
:: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
else {
asm volatile(
"{.reg .pred p;\n" \
"setp.eq.u32 p, 1, %4;\n" \
"tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, p;}\n"
:: "r"(d_tt_addr), "l"(a_desc), "l"(b_desc), "r"(idesc), "n"(acc)
);
}
}
}
template<int ncta=1> __device__ static inline void commit(kittens::semaphore &sem) {
if constexpr (ncta == 1) {
asm volatile(
"tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];\n"
:: "l"(&sem)
);
}
else {
asm volatile(
"tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;\n"
:: "l"(&sem), "h"((uint16_t)(0b11))
);
}
}
} // namespace tcgen05
} // namespace detail
template<typename T_AB> constexpr int reduction_dimension = sizeof(T_AB) == 2 ? 16 : sizeof(T_AB) == 4 ? 8 : 32; // haven't added fp4 yet.
// RS matmul equivalent
template<int trans_a, int n_trans_b, ducks::tt::all D, ducks::tt::all A, ducks::st_descriptor::input B, int acc=1, int ncta=1>
__device__ static inline void mma(D &d, const A &a, const B &b) {
constexpr int trans_b = 1 - n_trans_b;
// Do everything here.
constexpr int M = (trans_a ? A::cols : A::rows) * ncta;
static_assert(M == D::rows*ncta && ((ncta == 1 && (M == 64 || M == 128)) || (ncta == 2 && (M == 128 || M == 256)))); // output register is correctly sized
constexpr int N = (trans_b ? B::cols : B::rows) * ncta;
static_assert(N == D::cols); // output register is correctly sized
constexpr int K = trans_a ? A::rows : A::cols;
static_assert((trans_b ? B::rows : B::cols) == K); // K dimension must match
static_assert(std::is_same_v<typename A::T, typename B::T>); // A and B must match type.
// Usings
using T_AB = A::T; static_assert(std::is_same_v<T_AB, typename B::T>);
using T_D = D::T;
constexpr int red_dim = reduction_dimension<T_AB>;
static_assert(K%red_dim == 0, "K dimension must be divisible by red_dim.");
static_assert(
(std::is_same_v<T_D, half> && !std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && !std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && !std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, fp8e5m2>),
"Currently unsupported type combination for matrix multiply."
);
uint32_t idesc = detail::tcgen05::instruction_descriptor<T_D, T_AB, M, N, trans_a, trans_b, false>();
kittens::st_descriptor<ducks::st_descriptor::detail::get_st<B>, trans_b> b_desc(b);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
detail::tcgen05::template tt_st<T_AB, acc, ncta>(
d.addr,
a.template chunk_addr<trans_a>(0),
b_desc.chunk_descriptor(0),
idesc
);
#pragma unroll
for(int i = 1; i < K/red_dim; i++) {
detail::tcgen05::template tt_st<T_AB, 1, ncta>(
d.addr,
a.template chunk_addr<trans_a>(i),
b_desc.chunk_descriptor(i),
idesc
);
}
}
template<int trans_a, int n_trans_b, ducks::tt::all D, ducks::tt::all A, ducks::st_descriptor::input B, int acc=1, int ncta=1>
__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) {
mma<trans_a, n_trans_b, D, A, B, acc, ncta>(d, a, b);
detail::tcgen05::commit<ncta>(sem);
}
// SS matmul equivalent
template<int trans_a, int n_trans_b, ducks::tt::all D, ducks::st_descriptor::input A, ducks::st_descriptor::input B, int acc=1, int ncta=1>
__device__ static inline void mma(D &d, const A &a, const B &b) {
constexpr int trans_b = 1 - n_trans_b;
// Do everything here.
constexpr int M = (trans_a ? A::cols : A::rows) * ncta;
static_assert(M == D::rows*ncta && ((ncta == 1 && (M == 64 || M == 128)) || (ncta == 2 && (M == 128 || M == 256)))); // output register is correctly sized
constexpr int N = (trans_b ? B::cols : B::rows) * ncta;
static_assert(N == D::cols); // output register is correctly sized
constexpr int K = trans_a ? A::rows : A::cols;
static_assert((trans_b ? B::rows : B::cols) == K); // K dimension must match
static_assert(std::is_same_v<typename A::T, typename B::T>); // A and B must match type.
// Usings
using T_AB = A::T; static_assert(std::is_same_v<T_AB, typename B::T>);
using T_D = D::T;
constexpr int red_dim = reduction_dimension<T_AB>;
static_assert(K%red_dim == 0, "K dimension must be divisible by red_dim.");
static_assert(
(std::is_same_v<T_D, half> && !std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && !std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && !std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && !std::is_same_v<T_AB, fp8e5m2>),
"Currently unsupported type combination for matrix multiply."
);
uint32_t idesc = detail::tcgen05::instruction_descriptor<T_D, T_AB, M, N, trans_a, trans_b, false>();
kittens::st_descriptor<ducks::st_descriptor::detail::get_st<A>, trans_a> a_desc(a);
kittens::st_descriptor<ducks::st_descriptor::detail::get_st<B>, trans_b> b_desc(b);
asm volatile ("fence.proxy.async.shared::cta;\n" ::: "memory");
detail::tcgen05::template st_st<T_AB, acc, ncta>(
d.addr,
a_desc.chunk_descriptor(0),
b_desc.chunk_descriptor(0),
idesc
);
#pragma unroll
for(int i = 1; i < K/red_dim; i++) {
detail::tcgen05::template st_st<T_AB, 1, ncta>(
d.addr,
a_desc.chunk_descriptor(i),
b_desc.chunk_descriptor(i),
idesc
);
}
}
template<int trans_a, int n_trans_b, ducks::tt::all D, ducks::st_descriptor::input A, ducks::st_descriptor::input B, int acc=1, int ncta=1>
__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) {
mma<trans_a, n_trans_b, D, A, B, acc, ncta>(d, a, b);
detail::tcgen05::commit<ncta>(sem);
}
// Accumulator / numcta wrappers
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B, int acc=1>
__device__ static inline void mma2(D &d, const A &a, const B &b, semaphore &sem) {
mma<trans_a, trans_b, D, A, B, acc, 2>(d, a, b, sem);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B, int acc=1>
__device__ static inline void mma2(D &d, const A &a, const B &b) {
mma<trans_a, trans_b, D, A, B, acc, 2>(d, a, b);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm(D &d, const A &a, const B &b, semaphore &sem) {
mma<trans_a, trans_b, D, A, B, 0>(d, a, b, sem);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm(D &d, const A &a, const B &b) {
mma<trans_a, trans_b, D, A, B, 0>(d, a, b);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2(D &d, const A &a, const B &b, semaphore &sem) {
mma2<trans_a, trans_b, D, A, B, 0>(d, a, b, sem);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2(D &d, const A &a, const B &b) {
mma2<trans_a, trans_b, D, A, B, 0>(d, a, b);
}
// Transpose wrappers
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AB(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AB(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_ABt(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_ABt(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtB(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtB(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtBt(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AB(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AB(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_ABt(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::T, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_ABt(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::T, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtB(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtB(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtBt(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::T, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::T, D, A, B, 0>(d, a, b);
}
} // namespace kittens

13
tinygrad_repo/extra/thunder/cuda/include/ops/thread/thread.cuh Normal file
View File

@@ -0,0 +1,13 @@
/**
* @file
* @brief An aggregate header of all warp (worker) operations defined by ThunderKittens
*/
#pragma once
// no namespace wrapper needed here
#include "memory/memory.cuh"
#ifdef KITTENS_BLACKWELL
#include "mma/mma.cuh"
#endif