forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
IQ.Lvbs history cleanup
225
tinygrad_repo/extra/thunder/metal/include/common/utils.metal
Normal file
225
tinygrad_repo/extra/thunder/metal/include/common/utils.metal
Normal file
@@ -0,0 +1,225 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief General utilities for Thundermittens.
|
||||
*/
|
||||
#pragma once // not done
|
||||
/*
|
||||
TODO:
|
||||
shared allocator
|
||||
max shared mem for other hardware
|
||||
*/
|
||||
|
||||
#include <metal_stdlib>
|
||||
#include "base_types.metal"
|
||||
/**
|
||||
* @namespace mittens
|
||||
*
|
||||
* @brief The main namespace of Thundermittens.
|
||||
*/
|
||||
namespace mittens {
|
||||
/**
|
||||
* @namespace ore
|
||||
*
|
||||
* @brief The main namespace of Thundermittens Metal.
|
||||
*/
|
||||
|
||||
/* ---------- GENERAL CONSTANTS FOR mittens ---------- */
|
||||
|
||||
/**
|
||||
* @brief Tile dimension constant.
|
||||
*/
|
||||
constant constexpr const int TILE_DIM{8};
|
||||
constant constexpr const int TILE_ELEMENTS{TILE_DIM*TILE_DIM};
|
||||
constant constexpr const int SIMD_THREADS{32};
|
||||
|
||||
|
||||
#ifdef M2_PRO
|
||||
constant constexpr int MAX_SHARED_MEMORY = 32768;
|
||||
#else
|
||||
constant constexpr int MAX_SHARED_MEMORY = 32768;
|
||||
#endif
|
||||
/* ---------- TYPE HELPERS ---------- */
|
||||
/**
|
||||
* @namespace ducks
|
||||
*
|
||||
* @brief Thundermittens' namespace for template metaprogramming..
|
||||
*
|
||||
* This includes primarily dummy types and concept wrappers, along
|
||||
* with a few additional utilities.
|
||||
*/
|
||||
namespace ducks {
|
||||
|
||||
/**
|
||||
* @brief A type representing an empty default for a template.
|
||||
*/
|
||||
struct default_type {};
|
||||
|
||||
// This macro can't be done as a template, so it doesn't really have a location in mittens.
|
||||
#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
|
||||
|
||||
|
||||
}
|
||||
|
||||
/* ---------- SHUFFLE UTILS ---------- */
|
||||
/**
|
||||
* @brief Mask constant for all active threads in a warp.
|
||||
*/
|
||||
constant static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
|
||||
|
||||
template<typename T>
|
||||
static METAL_FUNC T shfl_sync(thread const T &f, const ushort laneid) {
|
||||
return metal::simd_shuffle(f, laneid);
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat shfl_sync<bfloat>(thread const bf16 &f, const ushort laneid) {
|
||||
// return as_type<bf16>(metal::simd_shuffle(*(thread half*)(&f), laneid));
|
||||
float f_val = (float)f;
|
||||
float shfl_val = metal::simd_shuffle(f_val, laneid);
|
||||
return (bf16)shfl_val;
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat2 shfl_sync<bfloat2>(thread const bf16_2 &f, const ushort laneid) {
|
||||
// return as_type<bf16_2>(metal::simd_shuffle(*(thread half2*)(&f), laneid));
|
||||
float2 f_val = (float2)f;
|
||||
float2 shfl_val = metal::simd_shuffle(f_val, laneid);
|
||||
return (bf16_2)shfl_val;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
static METAL_FUNC T shfl_down_fill_sync(thread const T &f, thread const T& fill_data, const ushort laneid) {
|
||||
return metal::simd_shuffle_and_fill_down(f, laneid, fill_data);
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat shfl_down_fill_sync<bfloat>(thread const bfloat &f, thread const bfloat &fill_data, const ushort laneid) {
|
||||
// return as_type<bf16>(metal::simd_shuffle_and_fill_down(*(thread half*)(&f), *(thread half*)(&fill_data), laneid));
|
||||
float f_val = (float)f;
|
||||
float fill_data_f = (float)fill_data;
|
||||
float shfl_val = metal::simd_shuffle_and_fill_down(f_val, fill_data_f, laneid);
|
||||
return (bf16)shfl_val;
|
||||
}
|
||||
template<>
|
||||
METAL_FUNC bfloat2 shfl_down_fill_sync<bfloat2>(thread const bfloat2 &f, thread const bfloat2 &fill_data, const ushort laneid) {
|
||||
// return as_type<bf16_2>(metal::simd_shuffle_and_fill_down(*(thread half2*)(&f), *(thread half2*)(&fill_data), laneid));
|
||||
float2 f_val = (float2)f;
|
||||
float2 fill_data_f = (float2)fill_data;
|
||||
float2 shfl_val = metal::simd_shuffle_and_fill_down(f_val, fill_data_f, laneid);
|
||||
return (bf16_2)shfl_val;
|
||||
}
|
||||
/**
|
||||
* @brief Perform a shuffle down operation on a packed type synchronously across a warp.
|
||||
* @tparam T The type of the value to be shuffled.
|
||||
* @param mask[in] The mask of active threads.
|
||||
* @param f[in] The value to be shuffled.
|
||||
* @param delta[in] The number of positions to shuffle down.
|
||||
* @return The result of the shuffle operation.
|
||||
*/
|
||||
template<typename T>
|
||||
static METAL_FUNC T shfl_down_sync(thread const T &f, int delta) {
|
||||
return metal::simd_shuffle_rotate_down(f, delta);
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat shfl_down_sync<bfloat>(thread const bf16 &f, int delta) {
|
||||
// return base_types::convertor<bf16, float>::convert(metal::simd_shuffle_rotate_down(base_types::convertor<float, bf16>::convert(f), delta));
|
||||
// return as_type<bf16>(metal::simd_shuffle_rotate_down(*(thread half*)(&f), delta));
|
||||
float f_val = (float)f;
|
||||
float shfl_val = metal::simd_shuffle_rotate_down(f_val, delta);
|
||||
return (bf16)shfl_val;
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat2 shfl_down_sync<bfloat2>(thread const bf16_2 &f, int delta) {
|
||||
// return as_type<bf16_2>(metal::simd_shuffle_rotate_down(*(thread const half2*)(&f), delta));
|
||||
// return base_types::convertor<bf16_2, float2>::convert(metal::simd_shuffle_rotate_down(base_types::convertor<float2, bf16_2>::convert(f), delta));
|
||||
|
||||
float2 f_val = (float2)f;
|
||||
float2 shfl_val = metal::simd_shuffle_rotate_down(f_val, delta);
|
||||
return (bf16_2)shfl_val;
|
||||
// return as_type<bf16_2>(metal::simd_shuffle_rotate_down(*(thread half2*)(&f), delta));
|
||||
}
|
||||
|
||||
|
||||
/* ---------- LOOP UNROLLING UTILS ---------- */
|
||||
|
||||
namespace meta {
|
||||
template <int Start, int End, int Stride, bool = (Start < End)>
|
||||
struct unroll_i_in_range {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F f, Args... args) {
|
||||
f(Start, args...);
|
||||
unroll_i_in_range<Start + Stride, End, Stride>::run(f, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <int Start, int End, int Stride>
|
||||
struct unroll_i_in_range<Start, End, Stride, false> {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F, Args...) {
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
template <int Start, int End, int Stride, bool = (Start < End)>
|
||||
struct unroll_i_j_in_range_inner {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F f, int outerIndex, Args... args) {
|
||||
f(outerIndex, Start, args...);
|
||||
unroll_i_j_in_range_inner<Start + Stride, End, Stride>::run(f, outerIndex, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <int Start, int End, int Stride>
|
||||
struct unroll_i_j_in_range_inner<Start, End, Stride, false> {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F, int, Args...) {
|
||||
}
|
||||
};
|
||||
|
||||
template <int StartOuter, int EndOuter, int StrideOuter,
|
||||
int StartInner, int EndInner, int StrideInner,
|
||||
bool = (StartOuter < EndOuter)>
|
||||
struct unroll_i_j_in_range {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F f, Args... args) {
|
||||
unroll_i_j_in_range_inner<StartInner, EndInner, StrideInner>::run(
|
||||
f, StartOuter, args...
|
||||
);
|
||||
unroll_i_j_in_range<
|
||||
StartOuter + StrideOuter, EndOuter, StrideOuter,
|
||||
StartInner, EndInner, StrideInner
|
||||
>::run(f, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <int StartOuter, int EndOuter, int StrideOuter,
|
||||
int StartInner, int EndInner, int StrideInner>
|
||||
struct unroll_i_j_in_range<StartOuter, EndOuter, StrideOuter,
|
||||
StartInner, EndInner, StrideInner, false> {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F, Args...) {
|
||||
}
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
|
||||
template <int N>
|
||||
struct ReadVector {
|
||||
float _[N];
|
||||
};
|
||||
|
||||
/* ---------- SHARED MEMORY UTILS ---------- */
|
||||
|
||||
#define mittens_ALIGN_AS(n) alignas(n)
|
||||
#define mittens_DEFAULT_ALIGN mittens_ALIGN_AS(16)
|
||||
|
||||
/**
|
||||
* @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
|
||||
*/
|
||||
struct mittens_DEFAULT_ALIGN alignment_dummy { int dummy; };
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user