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

IQ.Pilot Release Commit @ 0798119

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

24
tinygrad_repo/extra/thunder/metal/include/ops/group/group.metal Normal file
View File

@@ -0,0 +1,24 @@
/**
* @file
* @brief An aggregate header of all group (multi-warp) operations defined by Thundermittens
*/
#pragma once
#include "../../common/common.metal"
#include "../../types/types.metal"
#include "../warp/warp.metal" // several group memory ops rely on underlying warp-scope ops
namespace mittens {
template<int N_WARPS>
struct group {
constant static constexpr int GROUP_WARPS = N_WARPS; // This alias produces nice parallelism.
constant static constexpr int GROUP_THREADS = N_WARPS * mittens::SIMD_THREADS; // This alias produces nice parallelism.
static METAL_FUNC int simd_laneid(const unsigned threadIdx) { return threadIdx % mittens::SIMD_THREADS; }
static METAL_FUNC int laneid (const unsigned threadIdx) { return threadIdx % GROUP_THREADS; }
static METAL_FUNC int warpid (const unsigned threadIdx) { return laneid(threadIdx) / mittens::SIMD_THREADS; }
static METAL_FUNC int groupid (const unsigned threadIdx) { return threadIdx / GROUP_THREADS; }
#include "memory/memory.metal"
#include "shared/shared.metal"
};
}

2
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/memory.metal Normal file
View File

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

132
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/tile/global_to_register.metal Normal file
View File

@@ -0,0 +1,132 @@
/**
* @file
* @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
*/
/**
* @brief Collaboratively loads data from a source array into row-major layout tiles.
*
* @tparam RT The row-major layout tile type.
* @tparam U The data type of the source array.
* @param dst[out] The destination tile to load data into.
* @param src[in] The source array to load data from.
* @param row_stride[in] The stride in elements between rows in the source array.
*/
template<typename RT, typename GL>
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
load(thread RT &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename GL::dtype;
using U2 = typename base_types::packing<U>::packed_type;
const device U *src = (device U*)&_src.template get<RT>(idx);
const int row_stride = _src.row_stride();
int warp_laneid = threadIdx % 32;
const int row_offset = dst.rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
#pragma clang loop unroll(full)
for(int i = 0; i < dst.height; i++) {
int row = simd_y + i * RT::tile_size;
#pragma clang loop unroll(full)
for(int j = 0; j < dst.width; j++) {
int col = simd_x + j * RT::tile_size;
T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src[row * row_stride + col])));
dst.tiles[i][j].data.thread_elements()[0] = src2[0];
dst.tiles[i][j].data.thread_elements()[1] = src2[1];
}
}
}
template<typename RT, typename GL>
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
load(thread RT &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename GL::dtype;
using U2 = typename base_types::packing<U>::packed_type;
const device U *src = (device U*)&_src.template get<RT>(idx);
const int row_stride = _src.row_stride();
int warp_laneid = threadIdx % 32;
const int row_offset = dst.rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;;
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
#pragma clang loop unroll(full)
for(int i = 0; i < dst.height; i++) {
int row = simd_y + i * RT::tile_size;
#pragma clang loop unroll(full)
for(int j = 0; j < dst.width; j++) {
int col = simd_x + j * RT::tile_size;
T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src[row * row_stride + col])));
dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[row * row_stride + col]);
dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[(row + 1) * row_stride + col]);
}
}
}
/**
* @brief Collaboratively stores data from register tiles to a destination array in global memory with a row-major layout.
*
* @tparam RT The register tile type with a row-major layout.
* @tparam U The data type of the destination array.
* @param[out] dst The destination array in global memory to store data into.
* @param[in] src The source register tile to store data from.
* @param row_stride[in] The stride in elements between rows in the destination array.
*/
template<typename RT, typename GL>
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
store(thread GL &_dst, thread const RT &src, thread const coord &idx, const int threadIdx) {
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename GL::dtype;
using U2 = typename base_types::packing<U>::packed_type;
device U *dst = (device U*)&(_dst.template get<RT>(idx));
const int row_stride = _dst.row_stride();
int warp_laneid = simd_laneid(threadIdx);
const int row_offset = src.rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
#pragma clang loop unroll(full)
for(int i = 0; i < src.height; i++) {
int row = simd_y + i * RT::tile_size;
#pragma clang loop unroll(full)
for(int j = 0; j < src.width; j++) {
int col = simd_x + j * RT::tile_size;
U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]));
*(device U2*)(&dst[row*row_stride + col]) = src2;
}
}
}
template<typename RT, typename GL>
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
store(thread GL &_dst, thread const RT &src, thread const coord &idx, const int threadIdx) {
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename GL::dtype;
using U2 = typename base_types::packing<U>::packed_type;
device U *dst = (device U*)&(_dst.template get<RT>(idx));
const int row_stride = _dst.row_stride();
int warp_laneid = simd_laneid(threadIdx);
const int row_offset = src.rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
// const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
// const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
#pragma clang loop unroll(full)
for(int i = 0; i < src.height; i++) {
int row = simd_y + i * RT::tile_size;
#pragma clang loop unroll(full)
for(int j = 0; j < src.width; j++) {
int col = simd_x + j * RT::tile_size;
dst[row*row_stride + col] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
dst[(row + 1) * row_stride + col] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
}
}
}

144
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/tile/global_to_shared.metal Normal file
View File

@@ -0,0 +1,144 @@
/**
* @file
* @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory.
*/
//template<typename ST, typename U>
//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
//load(int i,
// threadgroup ST *dst, device U* src,
// thread const int& group_laneid,
// thread const int& memcpy_per_row,
// thread const int& elem_per_memcpy,
// thread const int& row_stride)
//{
// int idx = i * GROUP_THREADS + group_laneid;
// int row = idx / memcpy_per_row;
// int col = (idx*elem_per_memcpy) % ST::cols;
// if (row < ST::rows) {
// *(threadgroup float4*)(&(*dst)[{row, col}]) = *(device float4*)(&src[row*row_stride + col]);
// }
//}
template<typename ST, typename GL>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
int group_laneid = threadIdx % GROUP_THREADS;
using T = typename ST::T;
using U = typename GL::dtype;
device U *src = (device U*)&_src.template get<ST>(idx);
const int row_stride = _src.row_stride();
using read_vector = ReadVector<1>;
// we can handle this many rows each time we run a memcpy_async
constexpr const int elem_per_memcpy = sizeof(read_vector)/sizeof(typename ST::dtype);
constexpr const int memcpy_per_row = ST::cols / elem_per_memcpy;
int total_calls = ((ST::height * ST::width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
#pragma clang loop unroll(full)
for(int i = 0; i < total_calls; i++) {
int idx = i * GROUP_THREADS + group_laneid;
int row = idx / memcpy_per_row;
int col = (idx*elem_per_memcpy) % dst.cols;
if (row<dst.rows && col < dst.cols) {
*(threadgroup read_vector*)(&dst[{row, col}]) = *(device read_vector*)(&src[row*row_stride + col]);
// *(threadgroup float*)(&dst[{row, col}]) = 1.0f;
}
}
// dst[{0, 0}] = base_types::convertor<T, float>::convert(1.f);
// dst[{0, 0}] = total_calls;
// meta::unroll_i_in_range<0, total_calls, 1>::run(load<ST, typename GL::dtype>, &dst, src, group_laneid, memcpy_per_row, elem_per_memcpy, row_stride);
}
//template<typename ST, typename GL>
//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
//load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
// int group_laneid = threadIdx % GROUP_THREADS;
// int groupid = threadIdx / GROUP_THREADS;
// int laneid = threadIdx % SIMD_THREADS;
//
// using U = typename GL::dtype;
// device U *src = (device U*)&_src.template get<ST>(idx);
// const int row_stride = _src.row_stride();
//
// int elem_per_memcpy = sizeof(float)/sizeof(typename ST::dtype);
// int memcpy_per_row = ST::cols / elem_per_memcpy;
// int total_calls = ((ST::height * ST::width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
// /*
// 1x16 or 8 x 128
// */
// int offset = ST::num_elements / (GROUP_WARPS);
//// int offset = group_laneid
// #pragma clang loop unroll(full)
// for(int i = 0; i < total_calls; i++) {
// int idx = i * SIMD_THREADS + laneid;
//// int idx = i * () + group_laneid;
// int row = idx / memcpy_per_row;
// int col = (idx*elem_per_memcpy) % dst.cols;
// if (row<dst.rows) {
// *(threadgroup float*)(&dst[{row, col}]) = *(device float*)(&src[row*row_stride + col]);
// }
// }
//}
//
//template<typename ST, typename GL>
//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
//load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
// int warp_id = threadIdx / SIMD_THREADS;
// int lane_id = threadIdx % SIMD_THREADS;
//// int N_WARPS = /* number of warps in your group */;
//
// using U = typename GL::dtype;
// device U *src = (device U*)&_src.template get<ST>(idx);
// const int row_stride = _src.row_stride();
//
// int elem_per_memcpy = sizeof(float)/sizeof(typename ST::dtype);
// int memcpy_per_row = ST::cols / elem_per_memcpy;
// int total_memcpy_elems = (ST::height * ST::cols) / elem_per_memcpy;
// int elems_per_warp = (total_memcpy_elems + N_WARPS - 1) / N_WARPS; // Ceiling division
//
// int start_idx = warp_id * elems_per_warp;
// int end_idx = metal::min(start_idx + elems_per_warp, total_memcpy_elems);
//
// #pragma clang loop unroll(full)
// for (int idx = start_idx + lane_id; idx < end_idx; idx += SIMD_THREADS) {
// int row = idx / memcpy_per_row;
// int col = (idx % memcpy_per_row) * elem_per_memcpy;
// if (row < ST::height) {
// *(threadgroup float*)(&dst[{row, col}]) = *(device float*)(&src[row * row_stride + col]);
// }
// }
//}
template<typename ST, typename GL>
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
store(thread const GL &_dst, threadgroup const ST &src, thread const coord &idx, const int threadIdx) {
int group_laneid = threadIdx % GROUP_THREADS;
using U = typename GL::dtype;
device U *dst = (device U*)&_dst.template get<ST>(idx);
const int row_stride = _dst.row_stride();
using read_vector = ReadVector<1>;
// we can handle this many rows each time we run a memcpy_async
int elem_per_memcpy = sizeof(read_vector)/sizeof(typename ST::dtype); // float/float -> 1
int memcpy_per_row = ST::cols / elem_per_memcpy; // 240 memcpy per row
int total_calls = ((src.height * src.width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
#pragma clang loop unroll(full)
for(int i = 0; i < total_calls; i++) {
int idx = i * GROUP_THREADS + group_laneid;
int row = idx / memcpy_per_row;
int col = (idx*elem_per_memcpy) % src.cols;
if (row<src.rows && col < src.cols) {
*(device read_vector*)(&dst[row*row_stride + col]) = *(threadgroup read_vector*)(&src[{row, col}]);
// *(device float*)(&dst[row*row_stride + col]) = 1.f;
}
}
// dst[0] = src[{0,0}];
// dst[0] = total_calls;
// dst[0] = base_types::convertor<U, float>::convert(1);
}

152
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/tile/shared_to_register.metal Normal file
View File

@@ -0,0 +1,152 @@
/**
* @file
* @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
*/
/**
* @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
*
* @tparam RT The register tile type
* @tparam ST The shared tile type
* @param dst[out] The destination register tile.
* @param src[in] The source shared tile.
*/
template<typename RT, typename ST>
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
load(thread RT &dst, threadgroup const ST &src, const int threadIdx) {
constexpr int height = ST::height;
constexpr int warp_height = RT::height;
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename ST::dtype;
using U2 = typename base_types::packing<U>::packed_type;
int warp_laneid = simd_laneid(threadIdx);
const int row_offset = RT::rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
#pragma clang loop unroll(full)
for(int i = 0; i < dst.height; i++) {
int row = simd_y + i * mittens::TILE_DIM;
#pragma clang loop unroll(full)
for(int j = 0; j < dst.width; j++) {
int col = simd_x + j * mittens::TILE_DIM;
T2 src2 = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[{row, col}])));
dst.tiles[i][j].data.thread_elements()[0] = src2[0];
dst.tiles[i][j].data.thread_elements()[1] = src2[1];
}
}
}
template<typename RT, typename ST>
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
load(thread RT &dst, threadgroup const ST &src, const int threadIdx) {
constexpr int height = ST::height;
constexpr int warp_height = RT::height;
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename ST::dtype;
using U2 = typename base_types::packing<U>::packed_type;
int warp_laneid = simd_laneid(threadIdx);
const int row_offset = RT::rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
#pragma clang loop unroll(full)
for(int i = 0; i < dst.height; i++) {
#pragma clang loop unroll(full)
for(int j = 0; j < dst.width; j++) {
int row = simd_y + i * mittens::TILE_DIM;
int col = simd_x + j * mittens::TILE_DIM;
dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[{row + 0, col}]);
dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[{row + 1, col}]);
}
}
}
/**
* @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
*
* @tparam RT The register tile type
* @tparam ST The shared tile type
* @param dst[out] The destination shared tile.
* @param src[in] The source register tile.
*/
template<typename ST, typename RT>
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
store(threadgroup ST &dst, thread const RT &src, const int threadIdx) {
constexpr int height = ST::height;
constexpr int warp_height = RT::height;
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename ST::dtype;
using U2 = typename base_types::packing<U>::packed_type;
int warp_laneid = simd_laneid(threadIdx);
const int row_offset = RT::rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
#pragma clang loop unroll(full)
for(int i = 0; i < RT::height; i++) {
int row = simd_y + i * mittens::TILE_DIM;
#pragma clang loop unroll(full)
for(int j = 0; j < RT::width; j++) {
int col = simd_x + j * mittens::TILE_DIM;
U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0],
src.tiles[i][j].data.thread_elements()[1]));
*(threadgroup U2*)(&dst[{row, col}]) = src2;
}
}
}
template<typename ST, typename RT>
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
store(threadgroup ST &dst, thread const RT &src, const int threadIdx) {
constexpr int height = ST::height;
constexpr int warp_height = RT::height;
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
using T = typename RT::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U = typename ST::dtype;
using U2 = typename base_types::packing<U>::packed_type;
int warp_laneid = simd_laneid(threadIdx);
const int row_offset = RT::rows * warpid(threadIdx);
const short qid = warp_laneid / 4;
// const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
// const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
#pragma clang loop unroll(full)
for(int i = 0; i < RT::height; i++) {
#pragma clang loop unroll(full)
for(int j = 0; j < RT::width; j++) {
int row = simd_y + i * mittens::TILE_DIM;
int col = simd_x + j * mittens::TILE_DIM;
// U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0],
// src.tiles[i][j].data.thread_elements()[1]));
// *(threadgroup U2*)(&dst[{row, col}]) = src2;
dst[{row + 0, col}] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
dst[{row + 1, col}] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
}
}
}

8
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/tile/tile.metal Normal file
View File

@@ -0,0 +1,8 @@
/**
* @file
* @brief An aggregate header of group memory operations on tiles.
*/
#include "shared_to_register.metal"
#include "global_to_register.metal"
#include "global_to_shared.metal"

47
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/vec/global_to_register.metal Normal file
View File

@@ -0,0 +1,47 @@
/**
* @file
* @brief Functions for a warpgroup to collaboratively transfer data directly between global memory and registers and back.
*/
/**
* @brief Collaboratively loads data into register vectors from a source array in global memory.
*
* @tparam RV The register vector type.
* @tparam U The data type of the source array.
* @param[out] dst The destination register vector to load data into.
* @param[in] src The source array in global memory to load data from.
*/
template<typename RV, typename GL>
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
load(thread RV &dst, thread const GL &_src, thread coord idx, const int threadIdx) {
using T = typename RV::dtype;
using U = typename GL::dtype;
using U2 = typename base_types::packing<U>::packed_type;
using T2 = typename base_types::packing<T>::packed_type;
idx.c += warpid(threadIdx);
// Call warp level store
::mittens::load(dst, _src, idx, simd_laneid(threadIdx));
}
/**
* @brief Collaboratively stores data from register vectors to a destination array in global memory.
*
* @tparam RV The register vector type.
* @tparam U The data type of the destination array.
* @param[out] dst The destination array in global memory to store data into.
* @param[in] src The source register vector to store data from.
*/
template<typename RV, typename GL>
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
store(thread GL &_dst, thread const RV &src, thread coord idx, const int threadIdx) {
using T = typename RV::dtype;
// using U2 = typename base_types::packing<U>::packed_type;
using T2 = typename base_types::packing<T>::packed_type;
idx.c += warpid(threadIdx);
// Call warp level store
::mittens::store(_dst, src, idx, simd_laneid(threadIdx));
}

59
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/vec/global_to_shared.metal Normal file
View File

@@ -0,0 +1,59 @@
/**
* @file
* @brief Group (collaborative warp) ops for loading shared vectors from and storing to global memory.
*/
/**
* @brief Loads data from global memory into shared memory vector.
*
* This function loads data from a global memory location pointed to by `src` into a shared memory vector `dst`.
* It calculates the number of elements that can be transferred in one operation based on the size ratio of `float4` to the data type of `SV`.
* The function ensures coalesced memory access and efficient use of bandwidth by dividing the work among threads in a warp.
*
* @tparam SV Shared vector type, must satisfy ducks::sv::all concept.
* @param dst Reference to the shared vector where the data will be loaded.
* @param src Pointer to the global memory location from where the data will be loaded.
*/
template<typename SV, typename GL>
METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
load(threadgroup SV &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
using U = typename GL::dtype;
using read_vector = ReadVector<1>;
constexpr int elem_per_transfer = sizeof(read_vector) / sizeof(typename SV::dtype);
constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
device U *src = (device U*)&_src.template get<SV>(idx);
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < total_calls; i+=GROUP_THREADS) {
if(i * elem_per_transfer < dst.length)
*(threadgroup read_vector*)&dst[i*elem_per_transfer] = *(device read_vector*)&src[i*elem_per_transfer];
}
}
/**
* @brief Stores data from a shared memory vector to global memory.
*
* This function stores data from a shared memory vector `src` to a global memory location pointed to by `dst`.
* Similar to the load function, it calculates the number of elements that can be transferred in one operation based on the size ratio of `float4` to the data type of `SV`.
* The function ensures coalesced memory access and efficient use of bandwidth by dividing the work among threads in a warp.
*
* @tparam SV Shared vector type, must satisfy ducks::sv::all concept.
* @param dst Pointer to the global memory location where the data will be stored.
* @param src Reference to the shared vector from where the data will be stored.
*/
template<typename SV, typename GL>
METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
store(thread const GL &_dst, threadgroup const SV &src, thread const coord &idx, const int threadIdx) {
using read_vector = ReadVector<1>;
using U = typename GL::dtype;
constexpr int elem_per_transfer = sizeof(read_vector) / sizeof(typename SV::dtype);
constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
device U *dst = (device U*)&_dst.template get<SV>(idx);
metal::simdgroup_barrier(metal::mem_flags::mem_none);
#pragma clang loop unroll(full)
for(int i = laneid(threadIdx); i < total_calls; i+= GROUP_THREADS) {
if(i * elem_per_transfer < src.length)
*(device read_vector*)&dst[i*elem_per_transfer] = *(threadgroup read_vector*)&src[i*elem_per_transfer]; // lmao it's identical
}
}

60
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/vec/shared_to_register.metal Normal file
View File

@@ -0,0 +1,60 @@
/**
* @file
* @brief Functions for a group to collaboratively transfer data directly between shared memory and registers and back.
*/
/**
* @brief Collaboratively load data from a shared vector into register vectors split across a warpgroup.
*
* @tparam RV The register vector type
* @tparam SV The shared vector type
* @param dst[out] The destination register vector.
* @param src[in] The source shared vector.
*/
template<typename RV, typename SV>
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
load(thread RV &dst, threadgroup const SV &_src, const int threadIdx) {
using T = typename RV::dtype;
using U = typename SV::dtype;
using U2 = typename base_types::packing<U>::packed_type;
using T2 = typename base_types::packing<T>::packed_type;
static_assert(SV::length == RV::length*N_WARPS, "rv and sv dimensions do not match");// confirm size correct
// threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim> &src = subvec_inplace<RV::outer_dim, SV>(_src, warpid(threadIdx));
// threadgroup subvec &src = subvec_inplace<RV::outer_dim, SV>(_src, warpid(threadIdx));
unsigned warpId = warpid(threadIdx);
using subvec = typename SV::template subvec<RV::length>;
threadgroup subvec& src = *(threadgroup subvec*)(&_src[warpId *RV::length]);
::mittens::load<RV, subvec>(dst, src, simd_laneid(threadIdx)); // warp-level
}
/**
* @brief Collaboratively store data into a shared vector from register vectors split across a warpgroup.
*
* @tparam RV The register vector type
* @tparam SV The shared vector type
* @param dst[out] The destination shared vector.
* @param src[in] The source register vector.
*/
template<typename SV, typename RV>
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
store(threadgroup SV &_dst, thread const RV &src, const int threadIdx) {
using T = typename RV::dtype;
using U = typename SV::dtype;
using T2 = typename base_types::packing<T>::packed_type;
using U2 = typename base_types::packing<U>::packed_type;
static_assert(SV::length == RV::length*N_WARPS, "rv and sv dimensions do not match");// confirm size correct
// threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim> &dst = subvec_inplace<RV::outer_dim, SV>(_dst, warpid(threadIdx));
// ::mittens::store<threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim>, RV>(dst, src, simd_laneid(threadIdx)); // warp-level
unsigned warpId = warpid(threadIdx);
using subvec = typename SV::template subvec<RV::length>;
threadgroup subvec& dst = *(threadgroup subvec*)(&_dst[warpId * RV::length]);
::mittens::store(dst, src, simd_laneid(threadIdx)); // warp-level
}

8
tinygrad_repo/extra/thunder/metal/include/ops/group/memory/vec/vec.metal Normal file
View File

@@ -0,0 +1,8 @@
/**
* @file
* @brief An aggregate header of group memory operations on vectors.
*/
#include "shared_to_register.metal"
#include "global_to_register.metal"
#include "global_to_shared.metal"

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"