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

3
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/shared.metal Normal file
View File

@@ -0,0 +1,3 @@
#pragma once
#include "tile/tile.metal"
#include "vec/vec.metal"

59
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/tile/conversions.metal Normal file
View File

@@ -0,0 +1,59 @@
/**
* @file
* @brief Conversions between shared tile types.
*/
#pragma once // not done, add subtile
#include "../../../../common/common.metal"
#include "../../../../types/types.metal"
namespace mittens {
/* ---------- COPIES ---------- */
/**
* @brief Copies data from one shared memory tile to another, potentially with different data types and layouts.
*
* @tparam T The data type of the destination tile.
* @tparam U The data type of the source tile.
* @tparam _height The height of the tile.
* @tparam _width The width of the tile.
* @tparam L1 The layout of the destination tile.
* @tparam L2 The layout of the source tile.
* @param[out] dst The destination tile.
* @param[in] src The source tile.
*/
template<typename T, typename U, int _height, int _width>
static METAL_FUNC void copy(threadgroup st<T, _height, _width> &dst, threadgroup const st<U, _height, _width> &src, const ushort laneid) {
#pragma clang loop unroll(full)
for(int i = laneid; i < dst.num_elements; i+=mittens::SIMD_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst[{row, col}] = base_types::convertor<T, U>::convert(src[{row, col}]);
}
}
///* ---------- SUBTILE ---------- */
//
///**
//* @brief Returns a reference to a subtile of the given shared tile.
//*
//* @tparam subtile_height The height of the subtile.
//* @tparam subtile_width The width of the subtile.
//* @tparam ST The type of the input tile, which must satisfy the ducks::st::all concept.
//* @param src The input tile.
//* @param row_idx The row index of the subtile, in units of subtile_height*16 elements.
//* @param col_idx The col index of the subtile, in units of subtile_width*16 elements.
//* @return A reference to the subtile.
//*
//* @note The subtile {height, width} must evenly divide the tile {height, width}.
//*/
//template<int subtile_height, int subtile_width, ducks::st::all ST>
//__device__ inline typename ST::subtile<subtile_height, subtile_width> subtile_inplace(ST &src, int row_idx, int col_idx) {
// static_assert(ST::height % subtile_height == 0);
// static_assert(ST::width % subtile_width == 0);
// return typename ST::subtile<subtile_height, subtile_width>(
// &src[0], subtile_height*16*row_idx, subtile_width*16*col_idx
// );
//}
}

485
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/tile/maps.metal Normal file
View File

@@ -0,0 +1,485 @@
/**
* @file
* @brief Warp-scope maps on shared tiles.
*/
#pragma once
#include "../../../../common/common.metal"
#include "../../../../types/types.metal"
namespace mittens {
/* ---------- Uniform tile maps (independent of layout) ---------- */
/**
* @brief Performs a uniform unary operation on a tile.
*
* This function applies a given unary operation to each element of the source tile and stores the result in the destination tile.
* The operation is applied independently to each element, without considering its position or the values of neighboring elements.
*
* @tparam op The unary operation to be applied. Must be specialized to support operation on the data type of T.
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the unary operation is applied.
*/
template<typename op, typename ST> // T2, w, h can be inferred from dst as long as op is specialized
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
unary_map(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
#pragma clang loop unroll(full)
for(int i = laneid; i < ST::num_elements; i += SIMD_THREADS) {
dst.data[i] = op::template op<typename ST::dtype>(src.data[i]);
}
}
/**
* @brief Performs a uniform binary operation on a tile with a scalar parameter.
*
* This function applies a given binary operation to each element of the source tile and a scalar parameter, then stores the result in the destination tile.
* The operation is applied independently to each element, treating the scalar parameter as the second operand for each operation.
*
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the scalar parameter.
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the binary operation is applied.
* @param[in] param The scalar parameter to be used as the second operand in the binary operation.
*/
template<typename op, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
bin_map(threadgroup ST &dst, threadgroup const ST &src, thread const typename ST::dtype &param, const short laneid) {
#pragma clang loop unroll(full)
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
dst.data[i] = op::template op<typename ST::dtype>(src.data[i], param);
}
}
/**
* @brief Performs a uniform binary operation on two tiles.
*
* This function applies a given binary operation to corresponding elements of two source tiles and stores the result in the destination tile.
* The operation is applied independently to each pair of elements, without considering their positions or the values of neighboring elements.
*
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T.
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile to which the binary operation is applied.
* @param[in] rhs The second source tile to which the binary operation is applied.
*/
template<typename op, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
bin_map(threadgroup ST &dst, threadgroup const ST &lhs, threadgroup const ST &rhs, const ushort laneid) {
#pragma clang loop unroll(full)
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst.data[i] = op::template op<typename ST::dtype>(lhs.data[i], rhs.data[i]);
}
}
/**
* @brief Performs a row-wise binary operation on a tile with a vector.
*
* This function applies a given binary operation to each row of the source tile and the corresponding element of the source vector,
* then stores the result in the destination tile. The operation is applied independently to each row, using the vector element as
* the second operand for each element in the row.
*
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
* @tparam V The type of the vector. Must have the same data type as T.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the binary operation is applied.
* @param[in] vec The source vector containing the second operand for each row operation.
*/
template<typename op, typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>, void>::type
row_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const ushort laneid) {
static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
static_assert(SV::length == ST::rows, "Vector length must match the number of rows in the tile");
#pragma clang loop unroll(full)
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
int row = i/ST::cols, col = i%ST::cols;
dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[row]);
}
}
/**
* @brief Performs a column-wise binary operation on a tile with a vector.
*
* This function applies a given binary operation to each column of the source tile and the corresponding element of the source vector,
* then stores the result in the destination tile. The operation is applied independently to each column, using the vector element as
* the second operand for each element in the column.
*
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
* @tparam V The type of the vector. Must have the same data type as T.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the binary operation is applied.
* @param[in] vec The source vector containing the second operand for each column operation.
*/
template<typename op, typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const ushort laneid) {
static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
static_assert(SV::length == ST::cols, "Vector length must match the number of columns in the tile");
#pragma clang loop unroll(full)
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[col]);
}
}
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
// const maps
/**
* @brief Sets all elements of the destination tile to zero.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
zero(threadgroup ST &dst, const ushort laneid) {
unary_map<base_ops::zero, ST>(dst, dst, laneid);
}
/**
* @brief Sets all elements of the destination tile to one.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
one(threadgroup ST &dst, const ushort laneid) {
unary_map<base_ops::one, ST>(dst, dst, laneid);
}
/**
* @brief Sets all elements of the destination tile to positive infinity.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
pos_infty(threadgroup ST &dst, const ushort laneid) {
unary_map<base_ops::pos_infty, ST>(dst, dst, laneid);
}
/**
* @brief Sets all elements of the destination tile to negative infinity.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
neg_infty(threadgroup ST &dst, const ushort laneid) {
unary_map<base_ops::neg_infty, ST>(dst, dst, laneid);
}
// unary maps
/**
* @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the exponential function is applied.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
exp(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
unary_map<base_ops::exp, ST>(dst, src, laneid);
}
/**
* @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile, in base 2.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the exponential function is applied.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
exp2(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
unary_map<base_ops::exp2, ST>(dst, src, laneid);
}
/**
* @brief Applies the natural logarithm function to each element of the source tile and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the natural logarithm function is applied.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
log(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
unary_map<base_ops::log, ST>(dst, src, laneid);
}
/**
* @brief Applies the absolute function to each element of the source tile and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the absolute function is applied.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
abs(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
unary_map<base_ops::abs, ST>(dst, src, laneid);
}
/**
* @brief Applies the rectified linear unit function to each element of the source tile and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source tile to which the rectified linear unit function is applied.
*/
template<typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
relu(threadgroup ST &dst, const threadgroup ST &src, const ushort laneid) {
unary_map<base_ops::relu, ST>(dst, src, laneid);
}
/**
* @brief Copies the elements of the source tile to the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] src The source data to be copied.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
copy(threadgroup ST &dst, thread const U &src, const ushort laneid) {
bin_map<base_ops::copy2, ST>(dst, dst, src, laneid);
}
// uniform binary maps
/**
* @brief Finds the maximum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile.
* @param[in] rhs The second source data.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
max(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
bin_map<base_ops::max, ST>(dst, lhs, rhs, laneid);
}
/**
* @brief Finds the minimum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile.
* @param[in] rhs The second source data.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
min(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
bin_map<base_ops::min, ST>(dst, lhs, rhs, laneid);
}
/**
* @brief Adds each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile.
* @param[in] rhs The second source data.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
add(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
bin_map<base_ops::sum, ST>(dst, lhs, rhs, laneid);
}
/**
* @brief Subtracts each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile.
* @param[in] rhs The second source data.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
sub(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
bin_map<base_ops::sub, ST>(dst, lhs, rhs, laneid);
}
/**
* @brief Multiplies each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile.
* @param[in] rhs The second source data.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
mul(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
bin_map<base_ops::mul, ST>(dst, lhs, rhs, laneid);
}
/**
* @brief Divides each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
*
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
* @param[out] dst The destination tile where the results are stored.
* @param[in] lhs The first source tile.
* @param[in] rhs The second source data.
*/
template<typename ST, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
div(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
bin_map<base_ops::div, ST>(dst, lhs, rhs, laneid);
}
// Row and col maps
/**
* @brief Adds row values to each row of a tile.
*
* @tparam T Tile type.
* @tparam V Column vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the addition on.
* @param row_values[in] Column vector containing values to add to each row.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
add_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
row_map<base_ops::sum, ST, SV>(dst, src, row_values, laneid);
}
/**
* @brief Subtracts row values from each row of a tile.
*
* @tparam T Tile type.
* @tparam V Column vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the subtraction on.
* @param row_values[in] Column vector containing values to subtract from each row.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
sub_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
row_map<base_ops::sub, ST, SV>(dst, src, row_values, laneid);
}
/**
* @brief Multiplies each row of a tile by row values.
*
* @tparam T Tile type.
* @tparam V Column vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the multiplication on.
* @param row_values[in] Column vector containing values to multiply each row by.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
mul_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
row_map<base_ops::mul, ST, SV>(dst, src, row_values, laneid);
}
/**
* @brief Divides each row of a tile by row values.
*
* @tparam T Tile type.
* @tparam V Column vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the division on.
* @param row_values[in] Column vector containing values to divide each row by.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_tile<SV>(), void>::type
div_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
row_map<base_ops::div, ST, SV>(dst, src, row_values, laneid);
}
/**
* @brief Broadcast a vector into into a tile's rows.
*
* @tparam T Tile type.
* @tparam V Column vector type.
* @param dst[out] Destination tile where the result is stored.
* @param row_values[in] Column vector containing values to broadcast into rows.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
broadcast_row(threadgroup ST &dst, threadgroup const SV &row_values, const ushort laneid) {
row_map<base_ops::copy2, ST, SV>(dst, dst, row_values, laneid);
}
// col maps
/**
* @brief Adds column values to each column of a tile.
*
* @tparam T Tile type.
* @tparam V Row vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the addition on.
* @param col_values[in] Row vector containing values to add to each column.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
add_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
col_map<base_ops::sum, ST, SV>(dst, src, col_values, laneid);
}
/**
* @brief Subtracts column values from each column of a tile.
*
* @tparam T Tile type.
* @tparam V Row vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the subtraction on.
* @param col_values[in] Row vector containing values to subtract from each column.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
sub_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
col_map<base_ops::sub, ST, SV>(dst, src, col_values, laneid);
}
/**
* @brief Multiplies each column of a tile by column values.
*
* @tparam T Tile type.
* @tparam V Row vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the multiplication on.
* @param col_values[in] Row vector containing values to multiply each column by.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
mul_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
col_map<base_ops::mul, ST, SV>(dst, src, col_values, laneid);
}
/**
* @brief Divides each column of a tile by column values.
*
* @tparam T Tile type.
* @tparam V Row vector type.
* @param dst[out] Destination tile where the result is stored.
* @param src[in] Source tile to apply the division on.
* @param col_values[in] Row vector containing values to divide each column by.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
div_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
col_map<base_ops::div, ST, SV>(dst, src, col_values, laneid);
}
/**
* @brief Broadcast a vector into into a tile's columns.
*
* @tparam T Tile type.
* @tparam V Row vector type.
* @param dst[out] Destination tile where the result is stored.
* @param row_values[in] Row vector containing values to broadcast into cols.
*/
template<typename ST, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
broadcast_col(threadgroup ST &dst, threadgroup const SV &col_values, const ushort laneid) {
col_map<base_ops::copy2, ST, SV>(dst, dst, col_values, laneid);
}
}

295
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/tile/reductions.metal Normal file
View File

@@ -0,0 +1,295 @@
/**
* @file
* @brief Warp-scope reductions on shared tiles.
*/
#pragma once
#include "../../../../common/common.metal"
#include "../../../../types/types.metal"
namespace mittens {
/**
* 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, typename SV, typename ST, bool reset>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_reduce(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
using dtype = typename SV::dtype;
#pragma clang loop unroll(full)
for (int row = laneid; row < ST::rows; row += mittens::SIMD_THREADS) {
dtype accum = src[{row, 0}];
#pragma clang loop unroll(full)
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, typename SV, typename ST, bool reset>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_reduce(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
using dtype = typename SV::dtype;
#pragma clang loop unroll(full)
for (int col = laneid; col < src.cols; col += mittens::SIMD_THREADS) {
dtype accum = src[int2(0, col)];
#pragma clang loop unroll(full)
for (int row = 1; row < src.rows; row++) {
accum = op::template op<dtype>(accum, src[int2(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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_max(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
row_reduce<base_ops::max, SV, ST, true>(row_accum, src, row_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_min(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
row_reduce<base_ops::min, SV, ST, true>(row_accum, src, row_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_sum(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
row_reduce<base_ops::sum, SV, ST, true>(row_accum, src, row_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_prod(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
row_reduce<base_ops::mul, SV, ST, true>(row_accum, src, row_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_max(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
row_reduce<base_ops::max, SV, ST, false>(row_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_min(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
row_reduce<base_ops::min, SV, ST, false>(row_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_sum(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
row_reduce<base_ops::sum, SV, ST, false>(row_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
row_prod(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
row_reduce<base_ops::mul, SV, ST, false>(row_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_max(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
col_reduce<base_ops::max, SV, ST, true>(col_accum, src, col_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_min(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
col_reduce<base_ops::min, SV, ST, true>(col_accum, src, col_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_sum(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
col_reduce<base_ops::sum, SV, ST, true>(col_accum, src, col_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_prod(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
col_reduce<base_ops::mul, SV, ST, true>(col_accum, src, col_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_max(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
col_reduce<base_ops::max, SV, ST, false>(col_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_min(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
col_reduce<base_ops::min, SV, ST, false>(col_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_sum(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
col_reduce<base_ops::sum, SV, ST, false>(col_accum, src, src_accum, laneid);
}
/**
* @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<typename SV, typename ST>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
col_prod(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
col_reduce<base_ops::mul, SV, ST, false>(col_accum, src, src_accum, laneid);
}
}

4
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/tile/tile.metal Normal file
View File

@@ -0,0 +1,4 @@
#pragma once
#include "conversions.metal"
#include "maps.metal"
#include "reductions.metal"

60
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/vec/conversions.metal Normal file
View File

@@ -0,0 +1,60 @@
/**
* @file
* @brief Warp-scope conversions on shared vectors.
*/
#pragma once // done!
#include "../../../../common/common.metal"
#include "../../../../types/types.metal"
namespace mittens {
/**
* @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<typename SV1, typename SV2>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV1>() && ducks::is_shared_vector<SV2>(), void>::type
copy(threadgroup SV1 &dst, threadgroup const SV2 &src, const ushort laneid) {
static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length.");
#pragma clang loop unroll(full)
for(int i = laneid; i < dst.length; i+=SIMD_THREADS) {
dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
}
}
/* ---------- SUBVEC ---------- */
/**
* @brief Returns a reference to a subvec of a given shared vector
*
* @tparam subvec_tiles The length, in subtiles, of the subvec.
* @tparam SV The type of the input vector, which must satisfy the ducks::sv::all concept.
* @param src The input tile.
* @param vec_idx The index of the subtile, in units of subvec_tiles*16 elements.
* @return A reference to the subvec.
*
* @note The subvec length must evenly divide the vector length.
*/
template<int subvec_tiles, typename SV>
//using subvec = typename SV::template subvec<SV::length>;
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), threadgroup typename SV::template subvec<typename SV::dtype, subvec_tiles>&>::type
subvec_inplace(threadgroup SV &src, int vec_idx) {
return *(threadgroup typename SV::template subvec<typename SV::dtype, subvec_tiles>*)(&src[vec_idx*TILE_DIM*subvec_tiles]);
}
}

278
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/vec/maps.metal Normal file
View File

@@ -0,0 +1,278 @@
/**
* @file
* @brief Warp-scope maps on shared vectors.
*/
#pragma once
#include "../../../../common/common.metal"
#include "../../../../types/types.metal"
namespace mittens {
/**
* @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, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
unary_op(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
metal::simdgroup_barrier(metal::mem_flags::mem_none);
#pragma clang loop unroll(full)
for(int cur = laneid; cur < SV::length; cur+=SIMD_THREADS) {
dst[cur] = op::template op<typename SV::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, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
bin_op(threadgroup SV &dst, threadgroup const SV &lhs, threadgroup const SV &rhs, const ushort laneid) {
#pragma clang loop unroll(full)
for(int cur = laneid; cur < SV::length; cur+=SIMD_THREADS) {
dst[cur] = op::template op<typename SV::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, typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
bin_op(threadgroup SV &dst, threadgroup const SV &src, thread const typename SV::T &param, const ushort laneid) {
metal::simdgroup_barrier(metal::mem_flags::mem_none);
#pragma clang loop unroll(full)
for(int cur = laneid; cur < SV::length; cur+=SIMD_THREADS) {
dst[cur] = op::template op<typename SV::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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
zero(threadgroup SV &dst, const ushort laneid) {
unary_op<base_ops::zero, SV>(dst, dst, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
one(threadgroup SV &dst, const ushort laneid) {
unary_op<base_ops::one, SV>(dst, dst, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
pos_infty(threadgroup SV &dst, const ushort laneid) {
unary_op<base_ops::pos_infty, SV>(dst, dst, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
neg_infty(threadgroup SV &dst, const ushort laneid) {
unary_op<base_ops::neg_infty, SV>(dst, dst, laneid);
}
// ---- 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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
copy(threadgroup SV &dst, thread const U &src, const ushort laneid) {
bin_op<base_ops::copy2, SV>(dst, dst, src, laneid); // 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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
exp(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
unary_op<base_ops::exp, SV>(dst, src, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
exp2(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
unary_op<base_ops::exp2, SV>(dst, src, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
log(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
unary_op<base_ops::log, SV>(dst, src, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
abs(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
unary_op<base_ops::abs, SV>(dst, src, laneid);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
relu(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
unary_op<base_ops::relu, SV>(dst, src, laneid);
}
// ---- 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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
max(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
bin_op<base_ops::max, SV>(dst, lhs, rhs, laneid);
}
/**
* @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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
min(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
bin_op<base_ops::min, SV>(dst, lhs, rhs, laneid);
}
/**
* @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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
add(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
bin_op<base_ops::sum, SV>(dst, lhs, rhs, laneid);
}
/**
* @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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
sub(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
bin_op<base_ops::sub, SV>(dst, lhs, rhs, laneid);
}
/**
* @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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
mul(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
bin_op<base_ops::mul, SV>(dst, lhs, rhs, laneid);
}
/**
* @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<typename SV, typename U>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
div(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
bin_op<base_ops::div, SV>(dst, lhs, rhs, laneid);
}
}

268
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/vec/reductions.metal Normal file
View File

@@ -0,0 +1,268 @@
/**
* @file
* @brief Warp-scope maps on shared vectors.
*/
#pragma once
#include "../../../../common/common.metal"
#include "../../../../types/types.metal"
namespace mittens {
/**
* @brief Performs a reduction operation on elements of a shared memory vector within a warp.
*
* This function applies a specified operation to reduce the elements of a shared memory vector `src` to a single value.
* The result is stored in `accum`. If the `reset` parameter is true, the reduction includes an initial value `src_accum`.
* The reduction operation is performed in a warp-wide context, ensuring synchronization between threads in the warp.
*
* @tparam op The operation to perform on the elements. Must provide a static `op` method.
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
* @tparam reset A boolean flag indicating whether to include an initial value in the reduction.
* @param[out] accum The result of the reduction operation.
* @param[in] src The shared memory vector to reduce.
* @param[in] src_accum The initial value to include in the reduction if `reset` is false.
*/
template<typename op, typename SV, bool reset>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
reduce(thread typename SV::dtype &dst_accum, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
using T = typename SV::dtype;
{
T accum = src[0];
for (int i = 1; i < SV::length; i++) {
accum = op::template op<T>(accum, src[i]);
}
dst_accum = shfl_sync(accum, 0);
return;
}
//
T accum;
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = op::template op<T>(accum, src[i]);
}
if (src.length >= 32) {
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 1));
metal::simdgroup_barrier(metal::mem_flags::mem_none);
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 2));
metal::simdgroup_barrier(metal::mem_flags::mem_none);
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 4));
metal::simdgroup_barrier(metal::mem_flags::mem_none);
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 8));
metal::simdgroup_barrier(metal::mem_flags::mem_none);
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 16));
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 16));
} else if (src.length == 24) {
T shfl_val = shfl_down_sync<T>(accum, 1);
accum = op::template op<T>(accum, shfl_val);
shfl_val = shfl_down_sync<T>(accum, 2);
accum = op::template op<T>(accum, shfl_val);
shfl_val = shfl_down_sync<T>(accum, 4);
accum = op::template op<T>(accum, shfl_val);
shfl_val = shfl_down_sync<T>(accum, 8);
if (laneid < 16) {
accum = op::template op<T>(accum, shfl_val);
}
shfl_val = shfl_down_sync<T>(accum, 16);
accum = op::template op<T>(accum, shfl_val);
} else if (src.length == 16) {
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
} else if (src.length == 8) {
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
}
if (!reset) accum = op::template op<T>(accum, src_accum);
dst_accum = shfl_sync(accum, 0);
}
/* ---------- 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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
max(thread typename SV::dtype &max_val, threadgroup const SV &src, const ushort laneid) {
// reduce<base_ops::max, SV, true>(max_val, src, max_val, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::neg_infty();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::max::template op<T>(accum, src[i]);
}
max_val = (T)metal::simd_max((float)accum);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
min(thread typename SV::dtype &min_val, threadgroup const SV &src, const ushort laneid) {
// reduce<base_ops::min, SV, true>(min_val, src, min_val);
using T = typename SV::dtype;
T accum = base_types::constants<T>::pos_infty();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::min::template op<T>(accum, src[i]);
}
min_val = (T)metal::simd_min((float)accum);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
sum(thread typename SV::dtype &sum_val, threadgroup const SV &src, const ushort laneid) {
// reduce<base_ops::sum, SV, true>(sum_val, src, sum_val, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::zero();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::min::template op<T>(accum, src[i]);
}
sum_val = (T)metal::simd_sum((float)accum);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
prod(thread typename SV::dtype &prod_val, threadgroup const SV &src, const ushort laneid) {
// reduce<base_ops::mul, SV, true>(prod_val, src, prod_val, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::one();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::min::template op<T>(accum, src[i]);
}
prod_val = (T)metal::simd_product((float)accum);
}
// 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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
max(thread typename SV::dtype &max_val, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
// reduce<base_ops::max, SV, false>(max_val, src, src_accum, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::neg_infty();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::max::template op<T>(accum, src[i]);
}
max_val = (T)metal::simd_max((float)accum);
max_val = base_ops::max::template op<T>(max_val, src_accum);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
min(thread typename SV::dtype &min_val, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
// reduce<base_ops::min, SV, false>(min_val, src, src_accum, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::pos_infty();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::max::template op<T>(accum, src[i]);
}
min_val = (T)metal::simd_min((float)accum);
min_val = base_ops::max::template op<T>(min_val, src_accum);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
sum(thread typename SV::dtype &sum_val, threadgroup const SV &src, threadgroup const typename SV::dtype &src_accum, const ushort laneid) {
// reduce<base_ops::sum, SV, false>(sum_val, src, src_accum, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::zero();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::max::template op<T>(accum, src[i]);
}
sum_val = (T)metal::simd_sum((float)accum);
sum_val = base_ops::max::template op<T>(sum_val, src_accum);
}
/**
* @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<typename SV>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
prod(thread typename SV::dtype &prod_val, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
// reduce<base_ops::mul, SV, false>(prod_val, src, src_accum, laneid);
using T = typename SV::dtype;
T accum = base_types::constants<T>::one();
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
for(int i = laneid + 32; i < SV::length; i+=32) {
accum = base_ops::max::template op<T>(accum, src[i]);
}
prod_val = (T)metal::simd_product((float)accum);
prod_val = base_ops::max::template op<T>(prod_val, src_accum);
}
}

4
tinygrad_repo/extra/thunder/metal/include/ops/warp/shared/vec/vec.metal Normal file
View File

@@ -0,0 +1,4 @@
#pragma once
#include "conversions.metal"
#include "maps.metal"
#include "reductions.metal"