forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Prebuilt Release @ ab07000
This commit is contained in:
IQ.Lvbs history cleanup
32
tinygrad_repo/extra/thunder/amd/include/ops/group/group.cuh
Normal file
32
tinygrad_repo/extra/thunder/amd/include/ops/group/group.cuh
Normal file
@@ -0,0 +1,32 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of all group (multi-warp) operations defined by ThunderKittens
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../common/common.cuh"
|
||||
#include "../../types/types.cuh"
|
||||
#include "../warp/warp.cuh" // several group memory ops rely on underlying warp-scope ops
|
||||
|
||||
// A "warpgroup" is a special group of 4 consecutive warps defined by NVIDIA for certain SM_90+ operations.
|
||||
#define KITTENS_CHECK_WARPGROUP static_assert(N_WARPS==4, "PTX warpgroup (N_WARPS=4) function called from a non-warpgroup group.");
|
||||
|
||||
namespace kittens {
|
||||
/*
|
||||
This is meant to be used with a `using group_N = kittens::group<NUM_WORKERS>;` at the start of every kernel.
|
||||
*/
|
||||
template<int N_WARPS>
|
||||
struct group {
|
||||
static constexpr int GROUP_WARPS = N_WARPS; // This alias produces nice parallelism.
|
||||
static constexpr int GROUP_THREADS = N_WARPS * kittens::WARP_THREADS; // This alias produces nice parallelism.
|
||||
__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; }
|
||||
__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; }
|
||||
__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; }
|
||||
|
||||
#include "memory/memory.cuh"
|
||||
};
|
||||
|
||||
using warpgroup = group<4>; // special scope commonly used by SM_90 and later.
|
||||
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of colaborative group memory movement operations
|
||||
*/
|
||||
|
||||
#include "tile/tile.cuh"
|
||||
#include "vec/vec.cuh"
|
||||
@@ -0,0 +1,44 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory.
|
||||
*/
|
||||
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
|
||||
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
|
||||
kittens::load<axis, assume_aligned, ST, GL, COORD, GROUP_THREADS>(dst, src, idx);
|
||||
}
|
||||
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>> // default case
|
||||
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
|
||||
kittens::load<2, false, ST, GL, COORD, GROUP_THREADS>(dst, src, idx);
|
||||
}
|
||||
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
|
||||
__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
|
||||
kittens::store<axis, assume_aligned, ST, GL, COORD, GROUP_THREADS>(dst, src, idx);
|
||||
}
|
||||
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>> // default case
|
||||
__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
|
||||
kittens::store<2, false, ST, GL, COORD, GROUP_THREADS>(dst, src, idx);
|
||||
}
|
||||
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL>
|
||||
__device__ static inline void prefill_swizzled_offsets(ST &dst, const GL &src, uint32_t *swizzled_offsets) {
|
||||
kittens::prefill_swizzled_offsets<axis, assume_aligned, ST, GL, GROUP_THREADS>(dst, src, swizzled_offsets);
|
||||
}
|
||||
template<ducks::st::all ST, ducks::gl::all GL>
|
||||
__device__ static inline void prefill_swizzled_offsets(ST &dst, const GL &src, uint32_t *swizzled_offsets) {
|
||||
kittens::prefill_swizzled_offsets<2, false, ST, GL, GROUP_THREADS>(dst, src, swizzled_offsets);
|
||||
}
|
||||
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
|
||||
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx, const uint32_t *swizzled_offsets) {
|
||||
kittens::load<axis, assume_aligned, ST, GL, COORD, GROUP_THREADS>(dst, src, idx, swizzled_offsets);
|
||||
}
|
||||
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>> // default case
|
||||
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx, const uint32_t *swizzled_offsets) {
|
||||
kittens::load<2, false, ST, GL, COORD, GROUP_THREADS>(dst, src, idx, swizzled_offsets);
|
||||
}
|
||||
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
|
||||
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx, const uint32_t *__restrict__ swizzled_offsets, i32x4 srd, const void* base_ptr, uint32_t lds_base) {
|
||||
kittens::load<axis, assume_aligned, ST, GL, COORD, GROUP_THREADS>(dst, src, idx, swizzled_offsets, srd, base_ptr, lds_base);
|
||||
}
|
||||
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
|
||||
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx, const uint32_t *__restrict__ swizzled_offsets, i32x4 srd, const void* base_ptr, uint32_t lds_base) {
|
||||
kittens::load<2, false, ST, GL, COORD, GROUP_THREADS>(dst, src, idx, swizzled_offsets, srd, base_ptr, lds_base);
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of group memory operations on tiles.
|
||||
*/
|
||||
|
||||
#include "global_to_shared.cuh"
|
||||
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group (collaborative warp) ops for loading shared vectors from and storing to global memory.
|
||||
*/
|
||||
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>> // default case
|
||||
__device__ static inline void load(SV &dst, const GL &src, const COORD &idx) {
|
||||
kittens::load<SV, GL, COORD, GROUP_THREADS>(dst, src, idx);
|
||||
}
|
||||
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>> // default case
|
||||
__device__ static inline void store(const GL &dst, const SV &src, const COORD &idx) {
|
||||
kittens::store<SV, GL, COORD, GROUP_THREADS>(dst, src, idx);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of group memory operations on vectors.
|
||||
*/
|
||||
|
||||
#include "global_to_shared.cuh"
|
||||
Reference in New Issue
Block a user