superrob/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

7
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/shared.cuh Normal file
View File

@@ -0,0 +1,7 @@
/**
* @file
* @brief An aggregate header of group operations on data in shared memory
*/
#include "tile/tile.cuh"
#include "vec/vec.cuh"

16
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/tile/conversions.cuh Normal file
View File

@@ -0,0 +1,16 @@
/**
* @file
* @brief Group conversions between different shared memory tile types.
*/
/* ---------- COPIES ---------- */
template<ducks::st::all ST1, ducks::st::all ST2>
__device__ static inline void copy(ST1 &dst, const ST2 &src) {
static_assert(ST1::height == ST2::height && ST1::width == ST2::width, "Tiles must have the same height and width");
#pragma unroll
for(int i = laneid(); i < dst.num_elements; i+=GROUP_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst[{row, col}] = base_types::convertor<typename ST1::dtype, typename ST2::dtype>::convert(src[{row, col}]);
}
}

236
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/tile/maps.cuh Normal file
View File

@@ -0,0 +1,236 @@
/**
* @file
* @brief Group maps on shared tiles.
*/
template<typename op, ducks::st::all T> // T2, w, h can be inferred from dst as long as op is specialized
__device__ static inline void unary_map(T &dst, const T &src) {
#pragma unroll
for(int i = laneid(); i < dst.num_elements; i += GROUP_THREADS) {
dst.data[i] = op::template op<typename T::dtype>(src.data[i]);
}
}
template<typename op, ducks::st::all T>
__device__ static inline void bin_map(T &dst, const T &src, const typename T::dtype &param) {
#pragma unroll
for(int i = laneid(); i < dst.num_elements; i += GROUP_THREADS) {
dst.data[i] = op::template op<typename T::dtype>(src.data[i], param);
}
}
template<typename op, ducks::st::all T>
__device__ static inline void bin_map(T &dst, const T &lhs, const T &rhs) {
#pragma unroll
for(int i = laneid(); i < dst.num_elements; i += GROUP_THREADS) {
dst.data[i] = op::template op<typename T::dtype>(lhs.data[i], rhs.data[i]);
}
}
template<typename op, ducks::st::all T, ducks::sv::all V>
__device__ static inline void row_map(T &dst, const T &src, const V &vec) {
static_assert(std::is_same<typename T::dtype, typename V::dtype>::value, "Tile and vector must have the same data type");
static_assert(V::length == T::rows, "Vector length must match the number of rows in the tile");
#pragma unroll
for(int i = laneid(); i < dst.num_elements; i += GROUP_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst[{row, col}] = op::template op<typename T::dtype>(src[{row, col}], vec[row]);
}
}
template<typename op, ducks::st::all T, ducks::sv::all V>
__device__ static inline void col_map(T &dst, const T &src, const V &vec) {
static_assert(std::is_same<typename T::dtype, typename V::dtype>::value, "Tile and vector must have the same data type");
static_assert(V::length == T::cols, "Vector length must match the number of columns in the tile");
#pragma unroll
for(int i = laneid(); i < dst.num_elements; i += GROUP_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst[{row, col}] = op::template op<typename T::dtype>(src[{row, col}], vec[col]);
}
}
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
// All of the annoying qualifiers *should* be automatically inferred during compile-time.
// So, syntax should just be kittens::add_row(tile, colvec);
// const maps
template<ducks::st::all T>
__device__ static inline void zero(T &dst) {
unary_map<base_ops::zero, T>(dst, dst);
}
template<ducks::st::all T>
__device__ static inline void one(T &dst) {
unary_map<base_ops::one, T>(dst, dst);
}
template<ducks::st::all T>
__device__ static inline void pos_infty(T &dst) {
unary_map<base_ops::pos_infty, T>(dst, dst);
}
template<ducks::st::all T>
__device__ static inline void neg_infty(T &dst) {
unary_map<base_ops::neg_infty, T>(dst, dst);
}
// unary maps
template<ducks::st::all T>
__device__ static inline void exp(T &dst, const T &src) {
unary_map<base_ops::exp, T>(dst, src);
}
template<ducks::st::all T>
__device__ static inline void exp2(T &dst, const T &src) {
unary_map<base_ops::exp2, T>(dst, src);
}
template<ducks::st::all T>
__device__ static inline void log(T &dst, const T &src) {
unary_map<base_ops::log, T>(dst, src);
}
template<ducks::st::all T>
__device__ static inline void log2(T &dst, const T &src) {
unary_map<base_ops::log2, T>(dst, src);
}
template<ducks::st::all T>
__device__ static inline void abs(T &dst, const T &src) {
unary_map<base_ops::abs, T>(dst, src);
}
template<ducks::st::all T>
__device__ static inline void relu(T &dst, const T &src) {
unary_map<base_ops::relu, T>(dst, src);
}
template<ducks::st::all T, typename U>
__device__ static inline void copy(T &dst, const U &src) {
bin_map<base_ops::copy, T>(dst, src);
}
// uniform binary maps
template<ducks::st::all T, typename U>
__device__ static inline void max(T &dst, const T &lhs, const U &rhs) {
bin_map<base_ops::max, T>(dst, lhs, rhs);
}
template<ducks::st::all T, typename U>
__device__ static inline void min(T &dst, const T &lhs, const U &rhs) {
bin_map<base_ops::min, T>(dst, lhs, rhs);
}
template<ducks::st::all T, typename U>
__device__ static inline void add(T &dst, const T &lhs, const U &rhs) {
bin_map<base_ops::sum, T>(dst, lhs, rhs);
}
template<ducks::st::all T, typename U>
__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) {
bin_map<base_ops::sub, T>(dst, lhs, rhs);
}
template<ducks::st::all T, typename U>
__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) {
bin_map<base_ops::mul, T>(dst, lhs, rhs);
}
template<ducks::st::all T, typename U>
__device__ static inline void div(T &dst, const T &lhs, const U &rhs) {
bin_map<base_ops::div, T>(dst, lhs, rhs);
}
// Row and col maps
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void add_row(T &dst, const T &src, const V &row_values) {
row_map<base_ops::sum, T, V>(dst, src, row_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void sub_row(T &dst, const T &src, const V &row_values) {
row_map<base_ops::sub, T, V>(dst, src, row_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void mul_row(T &dst, const T &src, const V &row_values) {
row_map<base_ops::mul, T, V>(dst, src, row_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void div_row(T &dst, const T &src, const V &row_values) {
row_map<base_ops::div, T, V>(dst, src, row_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void broadcast_row(T &dst, const V &row_values) {
row_map<base_ops::copy2, T, V>(dst, dst, row_values);
}
// col maps
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void add_col(T &dst, const T &src, const V &col_values) {
col_map<base_ops::sum, T, V>(dst, src, col_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void sub_col(T &dst, const T &src, const V &col_values) {
col_map<base_ops::sub, T, V>(dst, src, col_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void mul_col(T &dst, const T &src, const V &col_values) {
col_map<base_ops::mul, T, V>(dst, src, col_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void div_col(T &dst, const T &src, const V &col_values) {
col_map<base_ops::div, T, V>(dst, src, col_values);
}
template<ducks::st::all T, ducks::sv::all V>
__device__ static inline void broadcast_col(T &dst, const V &col_values) {
col_map<base_ops::copy2, T, V>(dst, dst, col_values);
}
// Templated versions of each
template<int axis, ducks::st::all T, ducks::sv::all V>
__device__ static inline void add(T &dst, const T &src, const V &col_values) {
if constexpr (axis == axis::COL) add_col(dst, src, col_values);
else add_row(dst, src, col_values);
}
template<int axis, ducks::st::all T, ducks::sv::all V>
__device__ static inline void sub(T &dst, const T &src, const V &col_values) {
if constexpr (axis == axis::COL) sub_col(dst, src, col_values);
else sub_row(dst, src, col_values);
}
template<int axis, ducks::st::all T, ducks::sv::all V>
__device__ static inline void mul(T &dst, const T &src, const V &col_values) {
if constexpr (axis == axis::COL) mul_col(dst, src, col_values);
else mul_row(dst, src, col_values);
}
template<int axis, ducks::st::all T, ducks::sv::all V>
__device__ static inline void div(T &dst, const T &src, const V &col_values) {
if constexpr (axis == axis::COL) div_col(dst, src, col_values);
else div_row(dst, src, col_values);
}
template<int axis, ducks::st::all T, ducks::sv::all V>
__device__ static inline void broadcast(T &dst, const V &col_values) {
if constexpr (axis == axis::COL) broadcast_col(dst, col_values);
else broadcast_row(dst, col_values);
}

372
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/tile/reductions.cuh Normal file
View File

@@ -0,0 +1,372 @@
/**
* @file
* @brief Group reductions on shared tiles.
*/
/**
* Performs row-wise reduction on a matrix using a specified operation.
*
* @tparam op The operation to be applied for reduction.
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type with row layout.
* @param row_accum The accumulator where the result of the reduction is stored.
* @param src The source matrix on which to perform the reduction.
* @param src_accum The initial value of the accumulator, used when reset is false.
* @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
*/
template<typename op, ducks::sv::all V, ducks::st::all T, bool reset>
__device__ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) {
using dtype = typename V::dtype;
for (int row = laneid(); row < src.rows; row += GROUP_THREADS) {
dtype accum = src[{row, 0}];
#pragma unroll
for (int col = 1; col < src.cols; col++) {
accum = op::template op<dtype>(accum, src[{row, col}]);
}
if (reset) {
row_accum[row] = accum;
} else {
row_accum[row] = op::template op<dtype>(src_accum[row], accum);
}
}
}
/**
* Performs column-wise reduction on a matrix using a specified operation.
*
* @tparam op The operation to be applied for reduction.
* @tparam V The shared vector type for the column accumulator.
* @tparam T The shared matrix type with column layout.
* @param col_accum The accumulator where the result of the reduction is stored.
* @param src The source matrix on which to perform the reduction.
* @param src_accum The initial value of the accumulator, used when reset is false.
* @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
*/
template<typename op, ducks::sv::all V, ducks::st::all T, bool reset>
__device__ static inline void col_reduce(V &col_accum, const T &src, const V &src_accum) {
using dtype = typename V::dtype;
for (int col = laneid(); col < src.cols; col += GROUP_THREADS) {
dtype accum = src[{0, col}];
#pragma unroll
for (int row = 1; row < src.rows; row++) {
accum = op::template op<dtype>(accum, src[{row, col}]);
}
if (reset) {
col_accum[col] = accum;
} else {
col_accum[col] = op::template op<dtype>(src_accum[col], accum);
}
}
}
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
/**
* @brief Store the maximum of each row of the src shared matrix in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_max(V &row_accum, const T &src) {
row_reduce<base_ops::max, V, T, true>(row_accum, src, row_accum);
}
/**
* @brief Store the minimum of each row of the src shared matrix in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_min(V &row_accum, const T &src) {
row_reduce<base_ops::min, V, T, true>(row_accum, src, row_accum);
}
/**
* @brief Store the sum of each row of the src shared matrix in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_sum(V &row_accum, const T &src) {
row_reduce<base_ops::sum, V, T, true>(row_accum, src, row_accum);
}
/**
* @brief Store the product of each row of the src shared matrix in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_prod(V &row_accum, const T &src) {
row_reduce<base_ops::mul, V, T, true>(row_accum, src, row_accum);
}
/**
* @brief Store the maximum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_max(V &row_accum, const T &src, const V &src_accum) {
row_reduce<base_ops::max, V, T, false>(row_accum, src, src_accum);
}
/**
* @brief Store the minimum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_min(V &row_accum, const T &src, const V &src_accum) {
row_reduce<base_ops::min, V, T, false>(row_accum, src, src_accum);
}
/**
* @brief Store the sum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_sum(V &row_accum, const T &src, const V &src_accum) {
row_reduce<base_ops::sum, V, T, false>(row_accum, src, src_accum);
}
/**
* @brief Store the product of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] row_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void row_prod(V &row_accum, const T &src, const V &src_accum) {
row_reduce<base_ops::mul, V, T, false>(row_accum, src, src_accum);
}
/**
* @brief Store the maximum of each column of the src shared matrix in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_max(V &col_accum, const T &src) {
col_reduce<base_ops::max, V, T, true>(col_accum, src, col_accum);
}
/**
* @brief Store the minimum of each column of the src shared matrix in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_min(V &col_accum, const T &src) {
col_reduce<base_ops::min, V, T, true>(col_accum, src, col_accum);
}
/**
* @brief Store the sum of each column of the src shared matrix in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_sum(V &col_accum, const T &src) {
col_reduce<base_ops::sum, V, T, true>(col_accum, src, col_accum);
}
/**
* @brief Store the product of each column of the src shared matrix in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_prod(V &col_accum, const T &src) {
col_reduce<base_ops::mul, V, T, true>(col_accum, src, col_accum);
}
/**
* @brief Store the maximum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_max(V &col_accum, const T &src, const V &src_accum) {
col_reduce<base_ops::max, V, T, false>(col_accum, src, src_accum);
}
/**
* @brief Store the minimum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_min(V &col_accum, const T &src, const V &src_accum) {
col_reduce<base_ops::min, V, T, false>(col_accum, src, src_accum);
}
/**
* @brief Store the sum of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_sum(V &col_accum, const T &src, const V &src_accum) {
col_reduce<base_ops::sum, V, T, false>(col_accum, src, src_accum);
}
/**
* @brief Store the product of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
*
* @tparam V The shared vector type for the row accumulator.
* @tparam T The shared matrix type.
* @param[out] col_accum The accumulator where the result of the reduction is stored.
* @param[in] src The source matrix on which to perform the reduction.
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
*/
template<ducks::sv::all V, ducks::st::all T>
__device__ static inline void col_prod(V &col_accum, const T &src, const V &src_accum) {
col_reduce<base_ops::mul, V, T, false>(col_accum, src, src_accum);
}
// templated versions of each
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void max(V &dst, const T &src, const V &src_accum) {
if constexpr (ax == axis::COL) row_max(dst, src, src_accum);
else col_max(dst, src, src_accum);
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline auto max(const T &src, const V &src_accum) {
V dst;
if constexpr (ax == axis::COL) row_max(dst, src, src_accum);
else col_max(dst, src, src_accum);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void max(V &dst, const T &src) {
if constexpr (ax == axis::COL) row_max(dst, src);
else col_max(dst, src);
}
template<int ax, ducks::st::all T>
__device__ static inline auto max(const T &src) {
using V = std::conditional_t<ax==axis::COL, typename T::col_vec, typename T::row_vec>;
V dst;
if constexpr (ax == axis::COL) row_max(dst, src);
else col_max(dst, src);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void min(V &dst, const T &src, const V &src_accum) {
if constexpr (ax == axis::COL) row_min(dst, src, src_accum);
else col_min(dst, src, src_accum);
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline auto min(const T &src, const V &src_accum) {
V dst;
if constexpr (ax == axis::COL) row_min(dst, src, src_accum);
else col_min(dst, src, src_accum);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void min(V &dst, const T &src) {
if constexpr (ax == axis::COL) row_min(dst, src);
else col_min(dst, src);
}
template<int ax, ducks::st::all T>
__device__ static inline auto min(const T &src) {
using V = std::conditional_t<ax==axis::COL, typename T::col_vec, typename T::row_vec>;
V dst;
if constexpr (ax == axis::COL) row_min(dst, src);
else col_min(dst, src);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void sum(V &dst, const T &src, const V &src_accum) {
if constexpr (ax == axis::COL) row_sum(dst, src, src_accum);
else col_sum(dst, src, src_accum);
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline auto sum(const T &src, const V &src_accum) {
V dst;
if constexpr (ax == axis::COL) row_sum(dst, src, src_accum);
else col_sum(dst, src, src_accum);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void sum(V &dst, const T &src) {
if constexpr (ax == axis::COL) row_sum(dst, src);
else col_sum(dst, src);
}
template<int ax, ducks::st::all T>
__device__ static inline auto sum(const T &src) {
using V = std::conditional_t<ax==axis::COL, typename T::col_vec, typename T::row_vec>;
V dst;
if constexpr (ax == axis::COL) row_sum(dst, src);
else col_sum(dst, src);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void prod(V &dst, const T &src, const V &src_accum) {
if constexpr (ax == axis::COL) row_prod(dst, src, src_accum);
else col_prod(dst, src, src_accum);
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline auto prod(const T &src, const V &src_accum) {
V dst;
if constexpr (ax == axis::COL) row_prod(dst, src, src_accum);
else col_prod(dst, src, src_accum);
return dst;
}
template<int ax, ducks::sv::all V, ducks::st::all T>
__device__ static inline void prod(V &dst, const T &src) {
if constexpr (ax == axis::COL) row_prod(dst, src);
else col_prod(dst, src);
}
template<int ax, ducks::st::all T>
__device__ static inline auto prod(const T &src) {
using V = std::conditional_t<ax==axis::COL, typename T::col_vec, typename T::row_vec>;
V dst;
if constexpr (ax == axis::COL) row_prod(dst, src);
else col_prod(dst, src);
return dst;
}

37
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/tile/tile.cuh Normal file
View File

@@ -0,0 +1,37 @@
/**
* @file
* @brief An aggregate header for group operations on shared tiles.
*/
#include "conversions.cuh"
#include "maps.cuh"
#include "reductions.cuh"
template<ducks::st::all ST>
__device__ static inline bool hasnan(const ST &src) {
KITTENS_CHECK_WARP
bool nan_detected = false;
#pragma unroll
for(int i = laneid(); i < ST::num_elements; i+=GROUP_THREADS) {
if constexpr (std::is_same_v<typename ST::T, float>) {
if(isnan(src[i])) {
nan_detected = true;
}
}
else if constexpr (std::is_same_v<typename ST::T, bf16>) {
if(isnan(__bfloat162float(src[i]))) {
nan_detected = true;
}
}
else if constexpr (std::is_same_v<typename ST::T, half>) {
if(isnan(__half2float(src[i]))) {
nan_detected = true;
}
}
else {
static_assert(sizeof(typename ST::T) == 999, "Unsupported dtype");
}
}
// Ballot across the warp to see if any lane detected a nan
return (__ballot_sync(0xffffffff, nan_detected) != 0);
}

27
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/vec/conversions.cuh Normal file
View File

@@ -0,0 +1,27 @@
/**
* @file
* @brief Group conversions on shared vectors.
*/
/**
* @brief Copies data from one shared vector to another, converting data types if necessary.
*
* This function copies data from the source shared vector `src` to the destination shared vector `dst`.
* If the data types of `src` and `dst` are the same, it performs a direct memory copy. Otherwise, it
* converts each element from the source data type to the destination data type using the appropriate
* converter before copying.
*
* @tparam SV1 The type of the destination shared vector, must satisfy the ducks::sv::all concept.
* @tparam SV2 The type of the source shared vector, must satisfy the ducks::sv::all concept.
* @param[out] dst The destination shared vector.
* @param[in] src The source shared vector.
* @note The lengths of `src` and `dst` must be equal. This is enforced at compile time.
*/
template<ducks::sv::all SV1, ducks::sv::all SV2>
__device__ static inline void copy(SV1 &dst, const SV2 &src) {
static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length.");
#pragma unroll
for(int i = laneid(); i < dst.length; i+=GROUP_THREADS) {
dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
}
}

259
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/vec/maps.cuh Normal file
View File

@@ -0,0 +1,259 @@
/**
* @file
* @brief Group maps on shared vectors.
*/
/**
* @brief Applies a unary operation to each element of a shared memory vector.
*
* @tparam op Unary operation type.
* @tparam T Shared memory vector type.
* @param dst[out] Destination vector in which to store the result.
* @param src[in] Source vector to apply the unary operation.
*/
template<typename op, ducks::sv::all T>
__device__ static inline void unary_op(T &dst, const T &src) {
#pragma unroll
for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) {
dst[cur] = op::template op<typename T::dtype>(src[cur]);
}
}
/**
* @brief Perform a binary operation on two shared vectors.
*
* @tparam op The binary operation to perform.
* @tparam T The type of the vectors.
* @param dst[out] The destination vector where the result is stored.
* @param lhs[in] The left-hand side vector for the operation.
* @param rhs[in] The right-hand side vector for the operation.
*/
template<typename op, ducks::sv::all T>
__device__ static inline void bin_op(T &dst, const T &lhs, const T &rhs) {
#pragma unroll
for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) {
dst[cur] = op::template op<typename T::dtype>(lhs[cur], rhs[cur]);
}
}
/**
* @brief Perform a binary operation on a shared vector and a scalar.
*
* @tparam op The binary operation to perform.
* @tparam T The type of the vector.
* @param dst[out] The destination vector where the result is stored.
* @param src[in] The source vector for the operation.
* @param param[in] The scalar parameter for the operation.
*/
template<typename op, ducks::sv::all T>
__device__ static inline void bin_op(T &dst, const T &src, const typename T::dtype &param) {
#pragma unroll
for(auto cur = laneid(); cur < T::length; cur+=GROUP_THREADS) {
dst[cur] = op::template op<typename T::dtype>(src[cur], param);
}
}
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
// ---- const ops ----
/**
* @brief Sets all elements of a shared memory vector to zero.
*
* @tparam T Shared memory vector type.
* @param dst[out] Destination vector to be set to zero.
*/
template<ducks::sv::all T>
__device__ static inline void zero(T &dst) {
unary_op<base_ops::zero, T>(dst, dst);
}
/**
* @brief Sets all elements of a shared memory vector to one.
*
* @tparam T Shared memory vector type.
* @param dst[out] Destination vector to be set to one.
*/
template<ducks::sv::all T>
__device__ static inline void one(T &dst) {
unary_op<base_ops::one, T>(dst, dst);
}
/**
* @brief Sets all elements of a shared memory vector to positive infinity.
*
* @tparam T Shared memory vector type.
* @param dst[out] Destination vector to be set to positive infinity.
*/
template<ducks::sv::all T>
__device__ static inline void pos_infty(T &dst) {
unary_op<base_ops::pos_infty, T>(dst, dst);
}
/**
* @brief Sets all elements of a shared memory vector to negative infinity.
*
* @tparam T Shared memory vector type.
* @param dst[out] Destination vector to be set to negative infinity.
*/
template<ducks::sv::all T>
__device__ static inline void neg_infty(T &dst) {
unary_op<base_ops::neg_infty, T>(dst, dst);
}
// ---- unary ops ----
/**
* @brief Copies the elements from one shared vector to another.
*
* @tparam T Shared vector type.
* @tparam U Type of the source vector.
* @param dst[out] Destination vector where the elements will be copied to.
* @param src[in] Source vector to copy the elements from.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void copy(T &dst, const U &src) {
bin_op<base_ops::copy2, T>(dst, dst, src); // the second arg is ignored here.
}
/**
* @brief Applies the exponential function element-wise to a shared vector.
*
* @tparam T Shared vector type.
* @param dst[out] Destination vector where the exponential values will be stored.
* @param src[in] Source vector to apply the exponential function to.
*/
template<ducks::sv::all T>
__device__ static inline void exp(T &dst, const T &src) {
unary_op<base_ops::exp, T>(dst, src);
}
/**
* @brief Applies the exponential function element-wise to a shared vector, in base 2.
*
* @tparam T Shared vector type.
* @param dst[out] Destination vector where the exponential values will be stored.
* @param src[in] Source vector to apply the exponential function to.
*/
template<ducks::sv::all T>
__device__ static inline void exp2(T &dst, const T &src) {
unary_op<base_ops::exp2, T>(dst, src);
}
/**
* @brief Applies the natural logarithm function element-wise to a shared vector.
*
* @tparam T Shared vector type.
* @param dst[out] Destination vector where the logarithm values will be stored.
* @param src[in] Source vector to apply the logarithm function to.
*/
template<ducks::sv::all T>
__device__ static inline void log(T &dst, const T &src) {
unary_op<base_ops::log, T>(dst, src);
}
/**
* @brief Applies the logarithm base 2 function element-wise to a shared vector.
*
* @tparam T Shared vector type.
* @param dst[out] Destination vector where the logarithm base 2 values will be stored.
* @param src[in] Source vector to apply the logarithm base 2 function to.
*/
template<ducks::sv::all T>
__device__ static inline void log2(T &dst, const T &src) {
unary_op<base_ops::log2, T>(dst, src);
}
/**
* @brief Applies the absolute value function element-wise to a shared vector.
*
* @tparam T Shared vector type.
* @param dst[out] Destination vector where the absolute values will be stored.
* @param src[in] Source vector to apply the absolute value function to.
*/
template<ducks::sv::all T>
__device__ static inline void abs(T &dst, const T &src) {
unary_op<base_ops::abs, T>(dst, src);
}
/**
* @brief Applies the rectified linear unit (ReLU) function element-wise to a shared vector.
*
* @tparam T Shared vector type.
* @param dst[out] Destination vector where the ReLU values will be stored.
* @param src[in] Source vector to apply the ReLU function to.
*/
template<ducks::sv::all T>
__device__ static inline void relu(T &dst, const T &src) {
unary_op<base_ops::relu, T>(dst, src);
}
// ---- binary ops ----
/**
* @brief Computes the element-wise maximum of two shared vectors.
*
* @tparam T Shared vector type.
* @tparam U Type of the second vector.
* @param dst[out] Destination vector where the maximum values will be stored.
* @param lhs[in] First vector for the maximum operation.
* @param rhs[in] Second vector for the maximum operation.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void max(T &dst, const T &lhs, const U &rhs) {
bin_op<base_ops::max, T>(dst, lhs, rhs);
}
/**
* @brief Computes the element-wise minimum of two shared vectors.
*
* @tparam T Shared vector type.
* @tparam U Type of the second vector.
* @param dst[out] Destination vector where the minimum values will be stored.
* @param lhs[in] First vector for the minimum operation.
* @param rhs[in] Second vector for the minimum operation.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void min(T &dst, const T &lhs, const U &rhs) {
bin_op<base_ops::min, T>(dst, lhs, rhs);
}
/**
* @brief Computes the element-wise sum of two shared vectors.
*
* @tparam T Shared vector type.
* @tparam U Type of the second vector.
* @param dst[out] Destination vector where the sum values will be stored.
* @param lhs[in] First vector for the sum operation.
* @param rhs[in] Second vector for the sum operation.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void add(T &dst, const T &lhs, const U &rhs) {
bin_op<base_ops::sum, T>(dst, lhs, rhs);
}
/**
* @brief Computes the element-wise difference of two shared vectors.
*
* @tparam T Shared vector type.
* @tparam U Type of the second vector.
* @param dst[out] Destination vector where the difference values will be stored.
* @param lhs[in] First vector for the difference operation.
* @param rhs[in] Second vector for the difference operation.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void sub(T &dst, const T &lhs, const U &rhs) {
bin_op<base_ops::sub, T>(dst, lhs, rhs);
}
/**
* @brief Computes the element-wise product of two shared vectors.
*
* @tparam T Shared vector type.
* @tparam U Type of the second vector.
* @param dst[out] Destination vector where the product values will be stored.
* @param lhs[in] First vector for the product operation.
* @param rhs[in] Second vector for the product operation.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void mul(T &dst, const T &lhs, const U &rhs) {
bin_op<base_ops::mul, T>(dst, lhs, rhs);
}
/**
* @brief Computes the element-wise division of two shared vectors.
*
* @tparam T Shared vector type.
* @tparam U Type of the second vector.
* @param dst[out] Destination vector where the division values will be stored.
* @param lhs[in] First vector for the division operation.
* @param rhs[in] Second vector for the division operation.
*/
template<ducks::sv::all T, typename U>
__device__ static inline void div(T &dst, const T &lhs, const U &rhs) {
bin_op<base_ops::div, T>(dst, lhs, rhs);
}

193
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/vec/reductions.cuh Normal file
View File

@@ -0,0 +1,193 @@
/**
* @file
* @brief Group reductions on shared vectors.
*/
// The fastest way to do this, under most circumstances, is actually to just have each warp replicate it.
// This is not true for enormous shared vectors, but doing that efficiently actually requires some extra scratch shared memory.
// So, this is sufficient for the time being.
template<typename op, ducks::sv::all SV, bool reset>
__device__ static inline void reduce(typename SV::dtype &dst_accum, const SV &src, const typename SV::dtype &src_accum) {
if constexpr (GROUP_WARPS == 1) {
using T = SV::dtype;
int lane = laneid();
T accum;
if(lane < src.length) accum = src[lane]; // initialize a register accumulator
__syncwarp();
for(int i = lane+kittens::WARP_THREADS; i < src.length; i+=kittens::WARP_THREADS) {
accum = op::template op<T>(accum, src[i]);
}
__syncwarp();
// We can now reduce within the warp.
if constexpr (src.length > 16) {
accum = op::template op<T>(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 16));
__syncwarp();
}
accum = op::template op<T>(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 8));
__syncwarp();
accum = op::template op<T>(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 4));
__syncwarp();
accum = op::template op<T>(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 2));
__syncwarp();
accum = op::template op<T>(accum, packed_shfl_down_sync(kittens::MASK_ALL, accum, 1));
__syncwarp();
if constexpr (!reset) accum = op::template op<T>(accum, src_accum);
// broadcast to all threads in the warp.
dst_accum = packed_shfl_sync(kittens::MASK_ALL, accum, 0); // everyone takes from warp leader
}
else {
::kittens::group<1>::reduce<op, SV, reset>(dst_accum, src, src_accum);
}
}
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
/**
* @brief Finds the maximum element in a shared memory vector.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] max_val The maximum value found in the vector.
* @param[in] src The shared memory vector to find the maximum in.
*/
template<ducks::sv::all SV>
__device__ static inline void max(typename SV::dtype &max_val, const SV &src) {
reduce<base_ops::max, SV, true>(max_val, src, max_val);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype max(const SV &src) {
typename SV::dtype max_val;
reduce<base_ops::max, SV, true>(max_val, src, max_val);
return max_val;
}
/**
* @brief Finds the minimum element in a shared memory vector.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] min_val The minimum value found in the vector.
* @param[in] src The shared memory vector to find the minimum in.
*/
template<ducks::sv::all SV>
__device__ static inline void min(typename SV::dtype &min_val, const SV &src) {
reduce<base_ops::min, SV, true>(min_val, src, min_val);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype min(const SV &src) {
typename SV::dtype min_val;
reduce<base_ops::min, SV, true>(min_val, src, min_val);
return min_val;
}
/**
* @brief Calculates the sum of elements in a shared memory vector.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] sum_val The sum of the values in the vector.
* @param[in] src The shared memory vector to sum.
*/
template<ducks::sv::all SV>
__device__ static inline void sum(typename SV::dtype &sum_val, const SV &src) {
reduce<base_ops::sum, SV, true>(sum_val, src, sum_val);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype sum(const SV &src) {
typename SV::dtype sum_val;
reduce<base_ops::sum, SV, true>(sum_val, src, sum_val);
return sum_val;
}
/**
* @brief Calculates the product of elements in a shared memory vector.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] prod_val The product of the values in the vector.
* @param[in] src The shared memory vector to multiply.
*/
template<ducks::sv::all SV>
__device__ static inline void prod(typename SV::dtype &prod_val, const SV &src) {
reduce<base_ops::mul, SV, true>(prod_val, src, prod_val);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype prod(const SV &src) {
typename SV::dtype prod_val;
reduce<base_ops::mul, SV, true>(prod_val, src, prod_val);
return prod_val;
}
// Three operand versions.
/**
* @brief Finds the maximum element in a shared memory vector and accumulates it with src_accum.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] max_val The maximum value found in the vector, accumulated with src_accum.
* @param[in] src The shared memory vector to find the maximum in.
* @param[in] src_accum The initial value to accumulate with the maximum value found.
*/
template<ducks::sv::all SV>
__device__ static inline void max(typename SV::dtype &max_val, const SV &src, const typename SV::dtype &src_accum) {
reduce<base_ops::max, SV, false>(max_val, src, src_accum);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype max(const SV &src, const typename SV::dtype &src_accum) {
typename SV::dtype max_val;
reduce<base_ops::max, SV, false>(max_val, src, src_accum);
return max_val;
}
/**
* @brief Finds the minimum element in a shared memory vector and accumulates it with src_accum.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] min_val The minimum value found in the vector, accumulated with src_accum.
* @param[in] src The shared memory vector to find the minimum in.
* @param[in] src_accum The initial value to accumulate with the minimum value found.
*/
template<ducks::sv::all SV>
__device__ static inline void min(typename SV::dtype &min_val, const SV &src, const typename SV::dtype &src_accum) {
reduce<base_ops::min, SV, false>(min_val, src, src_accum);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype min(const SV &src, const typename SV::dtype &src_accum) {
typename SV::dtype min_val;
reduce<base_ops::min, SV, false>(min_val, src, src_accum);
return min_val;
}
/**
* @brief Calculates the sum of elements in a shared memory vector and accumulates it with src_accum.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] sum_val The sum of the values in the vector, accumulated with src_accum.
* @param[in] src The shared memory vector to sum.
* @param[in] src_accum The initial value to accumulate with the sum of the vector.
*/
template<ducks::sv::all SV>
__device__ static inline void sum(typename SV::dtype &sum_val, const SV &src, const typename SV::dtype &src_accum) {
reduce<base_ops::sum, SV, false>(sum_val, src, src_accum);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype sum(const SV &src, const typename SV::dtype &src_accum) {
typename SV::dtype sum_val;
reduce<base_ops::sum, SV, false>(sum_val, src, src_accum);
return sum_val;
}
/**
* @brief Calculates the product of elements in a shared memory vector and accumulates it with src_accum.
*
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @param[out] prod_val The product of the values in the vector, accumulated with src_accum.
* @param[in] src The shared memory vector to multiply.
* @param[in] src_accum The initial value to accumulate with the product of the vector.
*/
template<ducks::sv::all SV>
__device__ static inline void prod(typename SV::dtype &prod_val, const SV &src, const typename SV::dtype &src_accum) {
reduce<base_ops::mul, SV, false>(prod_val, src, src_accum);
}
template<ducks::sv::all SV>
__device__ static inline typename SV::dtype prod(const SV &src, const typename SV::dtype &src_accum) {
typename SV::dtype prod_val;
reduce<base_ops::mul, SV, false>(prod_val, src, src_accum);
return prod_val;
}

38
tinygrad_repo/extra/thunder/cuda/include/ops/group/shared/vec/vec.cuh Normal file
View File

@@ -0,0 +1,38 @@
/**
* @file
* @brief An aggregate header for group operations on shared vectors.
*/
#include "conversions.cuh"
#include "maps.cuh"
// no group vector reductions as they would require additional shared memory and synchronization, and those side effects just aren't worth it.
// warp vector reductions should be plenty fast in 99.9% of situations.
template<ducks::sv::all SV>
__device__ static inline bool hasnan(const SV &src) {
KITTENS_CHECK_WARP
bool nan_detected = false;
#pragma unroll
for(int i = laneid(); i < SV::length; i+=GROUP_THREADS) {
if constexpr (std::is_same_v<typename SV::T, float>) {
if(isnan(src[i])) {
nan_detected = true;
}
}
else if constexpr (std::is_same_v<typename SV::T, bf16>) {
if(isnan(__bfloat162float(src[i]))) {
nan_detected = true;
}
}
else if constexpr (std::is_same_v<typename SV::T, half>) {
if(isnan(__half2float(src[i]))) {
nan_detected = true;
}
}
else {
static_assert(sizeof(typename SV::T) == 999, "Unsupported dtype");
}
}
// Ballot across the warp to see if any lane detected a nan
return (__ballot_sync(0xffffffff, nan_detected) != 0);
}