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/group/shared/shared.metal Normal file
View File

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

27
tinygrad_repo/extra/thunder/metal/include/ops/group/shared/tile/conversions.metal Normal file
View File

@@ -0,0 +1,27 @@
/**
* @file
* @brief Group conversions between different shared memory tile types.
*/
/* ---------- 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 int threadIdx) {
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < dst.num_elements; i+=GROUP_THREADS) {
int row = i/dst.cols, col = i%dst.cols;
dst[{row, col}] = base_types::convertor<T, U>::convert(src[{row, col}]);
}
}

475
tinygrad_repo/extra/thunder/metal/include/ops/group/shared/tile/maps.metal Normal file
View File

@@ -0,0 +1,475 @@
/**
* @file
* @brief Group maps on shared tiles.
*/
/**
* @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 int threadIdx) {
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_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 int threadIdx) {
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_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 int threadIdx) {
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
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 int threadIdx) {
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(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
int row = i/dst.cols, col = i%dst.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 int threadIdx) {
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(threadIdx); i < dst.num_elements; i += GROUP_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 ---------- */
// All of the annoying qualifiers *should* be automatically inferred during compile-time.
// So, syntax should just be mittens::add_row(tile, colvec);
// 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 int threadIdx) {
unary_map<base_ops::zero, ST>(dst, dst, threadIdx);
}
/**
* @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 int threadIdx) {
unary_map<base_ops::one, ST>(dst, dst, threadIdx);
}
/**
* @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 int threadIdx) {
unary_map<base_ops::pos_infty, ST>(dst, dst, threadIdx);
}
/**
* @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 int threadIdx) {
unary_map<base_ops::neg_infty, ST>(dst, dst, threadIdx);
}
// 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 int threadIdx) {
unary_map<base_ops::exp, ST>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
unary_map<base_ops::exp2, ST>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
unary_map<base_ops::log, ST>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
unary_map<base_ops::abs, ST>(dst, src, threadIdx);
}
/**
* @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, threadgroup const ST &src, const int threadIdx) {
unary_map<base_ops::relu, ST>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
bin_map<base_ops::copy, ST>(dst, src, threadIdx);
}
// 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 int threadIdx) {
bin_map<base_ops::max, ST>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_map<base_ops::min, ST>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_map<base_ops::sum, ST>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_map<base_ops::sub, ST>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_map<base_ops::mul, ST>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_map<base_ops::div, ST>(dst, lhs, rhs, threadIdx);
}
// 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 int threadIdx) {
row_map<base_ops::sum, ST, SV>(dst, src, row_values, threadIdx);
}
/**
* @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 int threadIdx) {
row_map<base_ops::sub, ST, SV>(dst, src, row_values, threadIdx);
}
/**
* @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 int threadIdx) {
row_map<base_ops::mul, ST, SV>(dst, src, row_values, threadIdx);
}
/**
* @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_vector<SV>(), void>::type
div_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
row_map<base_ops::div, ST, SV>(dst, src, row_values, threadIdx);
}
/**
* @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 int threadIdx) {
row_map<base_ops::copy2, ST, SV>(dst, dst, row_values, threadIdx);
}
// 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 int threadIdx) {
col_map<base_ops::sum, ST, SV>(dst, src, col_values, threadIdx);
}
/**
* @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>() && ducks::is_shared_vector<SV>(), void>::type
sub_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
col_map<base_ops::sub, ST, SV>(dst, src, col_values, threadIdx);
}
/**
* @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 int threadIdx) {
col_map<base_ops::mul, ST, SV>(dst, src, col_values, threadIdx);
}
/**
* @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>() && ducks::is_shared_vector<SV>(), void>::type
div_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
col_map<base_ops::div, ST, SV>(dst, src, col_values, threadIdx);
}
/**
* @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 int threadIdx) {
col_map<base_ops::copy2, ST, SV>(dst, dst, col_values, threadIdx);
}

284
tinygrad_repo/extra/thunder/metal/include/ops/group/shared/tile/reductions.metal Normal file
View File

@@ -0,0 +1,284 @@
/**
* @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, 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 int threadIdx) {
using dtype = typename SV::dtype;
for (int row = laneid(threadIdx); row < src.rows; row += GROUP_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 int threadIdx) {
using dtype = typename SV::dtype;
for (int col = laneid(threadIdx); col < src.cols; col += GROUP_THREADS) {
dtype accum = src[{0, col}];
#pragma clang loop unroll(full)
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<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 int threadIdx) {
row_reduce<base_ops::max, SV, ST, true>(row_accum, src, row_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::min, SV, ST, true>(row_accum, src, row_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::sum, SV, ST, true>(row_accum, src, row_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::mul, SV, ST, true>(row_accum, src, row_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::max, SV, ST, false>(row_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::min, SV, ST, false>(row_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::sum, SV, ST, false>(row_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
row_reduce<base_ops::mul, SV, ST, false>(row_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::max, SV, ST, true>(col_accum, src, col_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::min, threadgroup SV, threadgroup ST, true>(col_accum, src, col_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::sum, SV, ST, true>(col_accum, src, col_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::mul, SV, ST, true>(col_accum, src, col_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::max, SV, ST, false>(col_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::min, SV, ST, false>(col_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::sum, SV, ST, false>(col_accum, src, src_accum, threadIdx);
}
/**
* @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 int threadIdx) {
col_reduce<base_ops::mul, SV, ST, false>(col_accum, src, src_accum, threadIdx);
}

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

@@ -0,0 +1,3 @@
#include "conversions.metal"
#include "maps.metal"
#include "reductions.metal"

29
tinygrad_repo/extra/thunder/metal/include/ops/group/shared/vec/conversions.metal Normal file
View File

@@ -0,0 +1,29 @@
/**
* @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<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 int threadIdx) {
static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length.");
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < dst.length; i+=GROUP_THREADS) {
dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
}
}

267
tinygrad_repo/extra/thunder/metal/include/ops/group/shared/vec/maps.metal Normal file
View File

@@ -0,0 +1,267 @@
/**
* @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, 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 int threadIdx) {
#pragma clang loop unroll(full)
for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_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 int threadIdx) {
#pragma clang loop unroll(full)
for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_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::dtype &param, const int threadIdx) {
#pragma clang loop unroll(full)
for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_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 int threadIdx) {
unary_op<base_ops::zero, SV>(dst, dst, threadIdx);
}
/**
* @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 int threadIdx) {
unary_op<base_ops::one, SV>(dst, dst, threadIdx);
}
/**
* @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 int threadIdx) {
unary_op<base_ops::pos_infty, SV>(dst, dst, threadIdx);
}
/**
* @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 int threadIdx) {
unary_op<base_ops::neg_infty, SV>(dst, dst, threadIdx);
}
// ---- 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 int threadIdx) {
bin_op<base_ops::copy2, SV>(dst, dst, src, threadIdx); // 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 int threadIdx) {
unary_op<base_ops::exp, SV>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
unary_op<base_ops::exp2, SV>(dst, src, threadIdx);
}
/**
* @brief Applies the natural logarithm 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 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 int threadIdx) {
unary_op<base_ops::log, SV>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
unary_op<base_ops::abs, SV>(dst, src, threadIdx);
}
/**
* @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 int threadIdx) {
unary_op<base_ops::relu, SV>(dst, src, threadIdx);
}
// ---- 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 int threadIdx) {
bin_op<base_ops::max, SV>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_op<base_ops::min, SV>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_op<base_ops::sum, SV>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_op<base_ops::sub, SV>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_op<base_ops::mul, SV>(dst, lhs, rhs, threadIdx);
}
/**
* @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 int threadIdx) {
bin_op<base_ops::div, SV>(dst, lhs, rhs, threadIdx);
}

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

@@ -0,0 +1,3 @@
#include "conversions.metal"
#include "maps.metal"