IQ.Pilot Release Commit @ b6534c0
This commit is contained in:
IQ.Lvbs CI [bot]
@@ -0,0 +1,392 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Basic operations on generic types.
|
||||
*/
|
||||
#pragma once
|
||||
#include "base_types.metal"
|
||||
#include <metal_math>
|
||||
|
||||
namespace mittens {
|
||||
/**
|
||||
* @namespace base_ops
|
||||
*
|
||||
* @brief A namespace for operations on basic data types.
|
||||
*/
|
||||
namespace base_ops {
|
||||
#define TEMPLATE_OPS_SINGLE(func_contents) \
|
||||
template<typename T> static METAL_FUNC T op(device const T &x) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &x) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &x) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_OVERRIDE_SINGLE(T, op_name, func_contents) \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &x) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &x) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &x) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_DOUBLE(func_contents) \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_OVERRIDE_DOUBLE(T, op_name, func_contents) \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_TRIPLE(func_contents) \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b, thread const T &c) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_OVERRIDE_TRIPLE(T, op_name, func_contents) \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b, thread const T &c) { func_contents }
|
||||
|
||||
|
||||
|
||||
/* ---------- CONST OPS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Represents the zero constant operation.
|
||||
*
|
||||
* This operation returns the zero value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the zero value.
|
||||
* @return The zero value of type T.
|
||||
*/
|
||||
struct zero {
|
||||
template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::zero(); }
|
||||
};
|
||||
/**
|
||||
* @brief Represents the one constant operation.
|
||||
*
|
||||
* This operation returns the one value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the one value.
|
||||
* @return The one value of type T.
|
||||
*/
|
||||
struct one {
|
||||
template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::one(); }
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Represents the positive infinity constant operation.
|
||||
*
|
||||
* This operation returns the positive infinity value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the positive infinity value.
|
||||
* @return The positive infinity value of type T.
|
||||
*/
|
||||
struct pos_infty {
|
||||
template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
|
||||
};
|
||||
/**
|
||||
* @brief Represents the negative infinity constant operation.
|
||||
*
|
||||
* This operation returns the negative infinity value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the negative infinity value.
|
||||
* @return The negative infinity value of type T.
|
||||
*/
|
||||
struct neg_infty {
|
||||
template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
|
||||
};
|
||||
|
||||
|
||||
/* ---------- UNARY OPS ---------- */
|
||||
/**
|
||||
* @brief Exponential function operation.
|
||||
*
|
||||
* This operation calculates the exponential of the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The exponential of the input value.
|
||||
*/
|
||||
struct exp {
|
||||
TEMPLATE_OPS_SINGLE(return metal::exp(x);)
|
||||
};
|
||||
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, exp, return bf16(metal::exp((float)x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, exp, return bf16_2(metal::exp(float2(x)));)
|
||||
|
||||
/**
|
||||
* @brief Exponential function operation, in base 2
|
||||
*
|
||||
* This operation calculates the exponential of the input value, in base 2.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The exponential of the input value.
|
||||
*/
|
||||
struct exp2 {
|
||||
template<typename T> static METAL_FUNC T op(device const T &x) { return metal::exp2(x); } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &x) { return metal::exp2(x); } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &x) { return metal::exp2(x); }
|
||||
};
|
||||
|
||||
//template<> METAL_FUNC bf16 exp2::op<bf16>(device const bf16 &x) { return bf16(metal::exp2(x)); } \
|
||||
//template<> METAL_FUNC bf16 exp2::op<bf16>(threadgroup const bf16 &x) { return bf16(metal::exp2(x)); } \
|
||||
//template<> METAL_FUNC bf16 exp2::op<bf16>(thread const bf16 &x) { return bf16(metal::exp2(x)); }
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, exp2, return bf16(metal::exp2(x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, exp2, return bf16_2(metal::exp2((float2)x));)
|
||||
|
||||
/**
|
||||
* @brief Natural log function operation.
|
||||
*
|
||||
* This operation calculates the natural logarithm of the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The natural logarithm of the input value.
|
||||
*/
|
||||
struct log {
|
||||
TEMPLATE_OPS_SINGLE(return metal::log(x);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, log, return bf16(metal::log(x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, log, return bf16_2(metal::log((float2)x));)
|
||||
|
||||
/**
|
||||
* @brief Absolute value operation.
|
||||
*
|
||||
* This operation calculates the absolute value of the input.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The absolute value of the input.
|
||||
*/
|
||||
struct abs {
|
||||
TEMPLATE_OPS_SINGLE(return metal::abs(x);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16 , abs, return bf16(metal::abs((float)x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, abs, return bf16_2(metal::abs((float2)x));)
|
||||
/**
|
||||
* @brief Rectified Linear Unit (ReLU) operation.
|
||||
*
|
||||
* This operation applies the ReLU function to the input, which is the
|
||||
* maximum of zero and the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The result of ReLU function applied to the input.
|
||||
*/
|
||||
struct relu {
|
||||
TEMPLATE_OPS_SINGLE(return max(x, base_types::constants<T>::zero());)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16 , relu, return bf16(metal::max((float)x, base_types::constants<float>::zero()));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, relu, return bf16_2(metal::max((float2)x, base_types::constants<float2>::zero()));)
|
||||
/**
|
||||
* @brief Copy operation.
|
||||
*
|
||||
* This operation returns the input value unchanged.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The input value.
|
||||
* @return The same value as the input.
|
||||
*/
|
||||
struct copy { // for non-compile-time setters.
|
||||
TEMPLATE_OPS_SINGLE(return x;)
|
||||
};
|
||||
|
||||
/* ---------- BINARY OPS ---------- */
|
||||
|
||||
|
||||
/**
|
||||
* @brief Copy2 operation.
|
||||
*
|
||||
* This operation returns the second input value unchanged.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value (ignored).
|
||||
* @param b[in] The second input value.
|
||||
* @return The same value as the second input.
|
||||
*/
|
||||
struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
|
||||
TEMPLATE_OPS_DOUBLE(return b;)
|
||||
};
|
||||
/**
|
||||
* @brief Sum operation.
|
||||
*
|
||||
* This operation calculates the sum of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The sum of the input values.
|
||||
*/
|
||||
struct sum {
|
||||
TEMPLATE_OPS_DOUBLE(return a+b;)
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Subtraction operation.
|
||||
*
|
||||
* This operation calculates the difference between two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The difference between the input values.
|
||||
*/
|
||||
struct sub {
|
||||
TEMPLATE_OPS_DOUBLE(return a-b;)
|
||||
};
|
||||
/**
|
||||
* @brief Multiplication operation.
|
||||
*
|
||||
* This operation calculates the product of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The product of the input values.
|
||||
*/
|
||||
struct mul {
|
||||
TEMPLATE_OPS_DOUBLE(return a*b;)
|
||||
};
|
||||
/**
|
||||
* @brief Division operation.
|
||||
*
|
||||
* This operation calculates the quotient of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The quotient of the input values.
|
||||
*/
|
||||
struct div {
|
||||
TEMPLATE_OPS_DOUBLE(return a/b;)
|
||||
};
|
||||
/**
|
||||
* @brief Maximum operation.
|
||||
*
|
||||
* This operation calculates the maximum of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The maximum of the input values.
|
||||
*/
|
||||
struct max {
|
||||
TEMPLATE_OPS_DOUBLE(return metal::max(a,b);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16 , max, return (bf16)metal::max((float)a, (float)b);)
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16_2, max, return (bf16_2)metal::max((float2)a, (float2)b);)
|
||||
/**
|
||||
* @brief Minimum operation.
|
||||
*
|
||||
* This operation calculates the minimum of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The minimum of the input values.
|
||||
*/
|
||||
struct min {
|
||||
TEMPLATE_OPS_DOUBLE(return metal::min(a,b);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16 , min, return (bf16)metal::min((float)a, (float)b);)
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16_2, min, return (bf16_2)metal::min((float2)a, (float2)b);)
|
||||
|
||||
|
||||
/* ---------- TERNARY OPS ---------- */
|
||||
/**
|
||||
* @brief Fused multiply-add operation A * B + C.
|
||||
*
|
||||
* This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @param c[in] The third input value to be added.
|
||||
* @return The result of the fused multiply-add operation.
|
||||
*/
|
||||
struct fma_AxBtC {
|
||||
TEMPLATE_OPS_TRIPLE(return sum::op<T>(mul::op<T>(a, b), c);)
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Fused multiply-add operation A * C + B.
|
||||
*
|
||||
* This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
|
||||
* This is particularly useful for attention mechanisms in neural networks.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The third input value to be added.
|
||||
* @param c[in] The second input value.
|
||||
* @return The result of the fused multiply-add operation.
|
||||
*/
|
||||
struct fma_AxCtB { // this is the one needed for attention
|
||||
TEMPLATE_OPS_TRIPLE(return sum::op<T>(mul::op<T>(a, c), b);)
|
||||
};
|
||||
|
||||
#undef TEMPLATE_OPS_SINGLE
|
||||
#undef TEMPLATE_OPS_OVERRIDE_SINGLE
|
||||
#undef TEMPLATE_OPS_DOUBLE
|
||||
#undef TEMPLATE_OPS_OVERRIDE_DOUBLE
|
||||
#undef TEMPLATE_OPS_TRIPLE
|
||||
#undef TEMPLATE_OPS_OVERRIDE_TRIPLE
|
||||
} // base_ops
|
||||
} // mittens
|
||||
@@ -0,0 +1,321 @@
|
||||
|
||||
#pragma once
|
||||
|
||||
namespace mittens {
|
||||
|
||||
using bf16 = bfloat;
|
||||
using bf16_2 = bfloat2;
|
||||
using bf16_4 = bfloat4;
|
||||
//using half_2 = half2;
|
||||
|
||||
namespace ducks {
|
||||
namespace base_types {
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT1() {
|
||||
return metal::is_same<typename T::dtype, float>::value ||
|
||||
metal::is_same<typename T::dtype, bf16 >::value ||
|
||||
metal::is_same<typename T::dtype, half>::value;
|
||||
}
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT2() {
|
||||
return metal::is_same<typename T::dtype, float2>::value ||
|
||||
metal::is_same<typename T::dtype, bf16_2>::value ||
|
||||
metal::is_same<typename T::dtype, half2>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT1Type() {
|
||||
return metal::is_same<T, float>::value ||
|
||||
metal::is_same<T, bf16 >::value ||
|
||||
metal::is_same<T, half>::value;
|
||||
}
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT2Type() {
|
||||
return metal::is_same<T, float2>::value ||
|
||||
metal::is_same<T, bf16_2>::value ||
|
||||
metal::is_same<T, half2>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT1Ptr() {
|
||||
return metal::is_same<T, device float*>::value ||
|
||||
metal::is_same<T, threadgroup float*>::value ||
|
||||
metal::is_same<T, thread float*>::value ||
|
||||
metal::is_same<T, device bf16*>::value ||
|
||||
metal::is_same<T, threadgroup bf16*>::value ||
|
||||
metal::is_same<T, thread bf16*>::value ||
|
||||
metal::is_same<T, device half*>::value ||
|
||||
metal::is_same<T, threadgroup half*>::value ||
|
||||
metal::is_same<T, thread half*>::value;
|
||||
}
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT2Ptr() {
|
||||
return metal::is_same<T, device float2*>::value ||
|
||||
metal::is_same<T, threadgroup float2*>::value ||
|
||||
metal::is_same<T, thread float2*>::value ||
|
||||
metal::is_same<T, device bf16_2*>::value ||
|
||||
metal::is_same<T, threadgroup bf16_2*>::value ||
|
||||
metal::is_same<T, thread bf16_2*>::value ||
|
||||
metal::is_same<T, device half2*>::value ||
|
||||
metal::is_same<T, threadgroup half2*>::value ||
|
||||
metal::is_same<T, thread half2*>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isTKType() { // good enough
|
||||
return !isT1Type<T>() && !isT2Type<T>() && !isT1Ptr<T>() && !isT2Ptr<T>();
|
||||
}
|
||||
|
||||
} // namespace base_types
|
||||
} // namespace ducks
|
||||
|
||||
/**
|
||||
* @namespace base_types
|
||||
*
|
||||
* @brief A namespace for Thundermittens basic data types.
|
||||
*/
|
||||
namespace base_types {
|
||||
/**
|
||||
* @brief Provides compile-time constants for different types.
|
||||
*
|
||||
* @tparam T The type for which to provide constants.
|
||||
*/
|
||||
template<typename T> struct constants {
|
||||
/**
|
||||
* @brief Zero
|
||||
* @return Constexpr zero with type T
|
||||
*/
|
||||
static METAL_FUNC constexpr T zero() { return T{0}; }
|
||||
/**
|
||||
* @brief One
|
||||
* @return Constexpr one with type T
|
||||
*/
|
||||
static METAL_FUNC constexpr T one() { return T{1}; }
|
||||
/**
|
||||
* @brief Positive infinity. Particularly useful for initializing before a min op.
|
||||
* @return Constexpr positive infinity with type T
|
||||
*/
|
||||
static METAL_FUNC constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
|
||||
/**
|
||||
* @brief Negative infinity. Particularly useful for initializing before a max op.
|
||||
* @return Constexpr negative infinity with type T
|
||||
*/
|
||||
static METAL_FUNC constexpr T neg_infty() { return T{-INFINITY}; }
|
||||
};
|
||||
template<> struct constants<float> {
|
||||
static METAL_FUNC constexpr float zero() { return 0.f; }
|
||||
static METAL_FUNC constexpr float one() { return 1.f; }
|
||||
static METAL_FUNC constexpr float pos_infty() { return INFINITY; }
|
||||
static METAL_FUNC constexpr float neg_infty() { return -INFINITY; }
|
||||
};
|
||||
template<> struct constants<float2> {
|
||||
static METAL_FUNC constexpr float2 zero() { return float2(0.f, 0.f); }
|
||||
static METAL_FUNC constexpr float2 one() { return float2(1.f, 1.f); }
|
||||
static METAL_FUNC constexpr float2 pos_infty() { return float2(constants<float>::pos_infty(), constants<float>::pos_infty()); }
|
||||
static METAL_FUNC constexpr float2 neg_infty() { return float2(constants<float>::neg_infty(), constants<float>::neg_infty()); }
|
||||
};
|
||||
template<> struct constants<bf16> {
|
||||
static METAL_FUNC constexpr bf16 zero() { return 0.bf; }
|
||||
static METAL_FUNC constexpr bf16 one() { return 1.bf; }
|
||||
static METAL_FUNC constexpr bf16 pos_infty() { return HUGE_VALBF; }
|
||||
static METAL_FUNC constexpr bf16 neg_infty() { return -HUGE_VALBF; }
|
||||
};
|
||||
template<> struct constants<bf16_2> {
|
||||
static METAL_FUNC constexpr bf16_2 zero() { return bf16_2(constants<bf16>::zero(), constants<bf16>::zero()); }
|
||||
static METAL_FUNC constexpr bf16_2 one() { return bf16_2(constants<bf16>::one(), constants<bf16>::one()); }
|
||||
static METAL_FUNC constexpr bf16_2 pos_infty() { return bf16_2(constants<bf16>::pos_infty(), constants<bf16>::pos_infty()); }
|
||||
static METAL_FUNC constexpr bf16_2 neg_infty() { return bf16_2(constants<bf16>::neg_infty(), constants<bf16>::neg_infty()); }
|
||||
};
|
||||
template<> struct constants<half> {
|
||||
static METAL_FUNC constexpr half zero() { return half(0.h); }
|
||||
static METAL_FUNC constexpr half one() { return half(1.h); }
|
||||
static METAL_FUNC constexpr half pos_infty() { return HUGE_VALH; }
|
||||
static METAL_FUNC constexpr half neg_infty() { return -HUGE_VALH; }
|
||||
};
|
||||
|
||||
template<> struct constants<half2> {
|
||||
static METAL_FUNC constexpr half2 zero() { return half2(constants<half>::zero(), constants<half>::zero()); }
|
||||
static METAL_FUNC constexpr half2 one() { return half2(constants<half>::one(), constants<half>::one()); }
|
||||
static METAL_FUNC constexpr half2 pos_infty() { return half2(constants<half>::pos_infty(), constants<half>::pos_infty()); }
|
||||
static METAL_FUNC constexpr half2 neg_infty() { return half2(constants<half>::neg_infty(), constants<half>::neg_infty()); }
|
||||
};
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* @brief Provides information about packing of elements for a given type.
|
||||
*
|
||||
* @tparam T The type for which to provide packing information.
|
||||
*/
|
||||
template<typename T> struct packing {
|
||||
// /**
|
||||
// * @brief The number of elements packed together.
|
||||
// *
|
||||
// * @return constexpr int representing number of elements within the type.
|
||||
// */
|
||||
// static METAL_FUNC constexpr int num() { return 1; }
|
||||
// /**
|
||||
// * @brief Packs a single T element twice (replicated) into its packed type.
|
||||
// *
|
||||
// * @param i[in] The element to pack.
|
||||
// * @return The packed type.
|
||||
// */
|
||||
// static METAL_FUNC constexpr T pack(device const bf16 &i);
|
||||
// static METAL_FUNC constexpr T pack(threadgroup const bf16 &i);
|
||||
// static METAL_FUNC constexpr T pack(thread const bf16 &i);
|
||||
};
|
||||
|
||||
#define PACK_FUNCTIONS(T1, T2) \
|
||||
static METAL_FUNC constexpr T2 pack(device const T1 &i) { return T2{i, i}; } \
|
||||
static METAL_FUNC constexpr T2 pack(threadgroup const T1 &i) { return T2{i, i}; } \
|
||||
static METAL_FUNC constexpr T2 pack(thread const T1 &i) { return T2{i, i}; }
|
||||
|
||||
template<> struct packing<bf16> {
|
||||
static METAL_FUNC constexpr int num() { return 1; }
|
||||
using unpacked_type = bf16;
|
||||
using packed_type = bf16_2;
|
||||
using packed_four = bf16_4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<half> {
|
||||
static METAL_FUNC constexpr int num() { return 1; }
|
||||
using unpacked_type = half;
|
||||
using packed_type = half2;
|
||||
using packed_four = half4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<float> {
|
||||
static METAL_FUNC constexpr int num() { return 1; }
|
||||
using unpacked_type = float;
|
||||
using packed_type = float2;
|
||||
using packed_four = float4;
|
||||
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<bf16_2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
using unpacked_type = bf16;
|
||||
using packed_type = bf16_2;
|
||||
using packed_four = bf16_4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<half2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
using unpacked_type = half;
|
||||
using packed_type = half2;
|
||||
using packed_four = half4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<float2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
using unpacked_type = float;
|
||||
using packed_type = float2;
|
||||
using packed_four = float4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<int2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
};
|
||||
template<> struct packing<float4> {
|
||||
static METAL_FUNC constexpr int num() { return 4; }
|
||||
};
|
||||
template<> struct packing<int4> {
|
||||
static METAL_FUNC constexpr int num() { return 4; }
|
||||
};
|
||||
|
||||
|
||||
/**
|
||||
* @brief Provides templated functionality to convert between different types.
|
||||
*
|
||||
* @tparam T The target type for conversion.
|
||||
* @tparam U The source type for conversion.
|
||||
*/
|
||||
template<typename T, typename U> struct convertor {
|
||||
/**
|
||||
* @brief Converts a value of type U to type T.
|
||||
*
|
||||
* @param u[in] The value of type U to convert.
|
||||
* @return T The converted value of type T.
|
||||
*/
|
||||
static METAL_FUNC T convert(device const U & u) { return (T)u; }
|
||||
static METAL_FUNC T convert(threadgroup const U & u) { return (T)u; }
|
||||
static METAL_FUNC T convert(thread const U & u) { return (T)u; }
|
||||
};
|
||||
|
||||
template<> struct convertor<float, bf16> {
|
||||
// fptrunc float %_ to bfloat
|
||||
static METAL_FUNC float convert(device const bf16 & u) { return float(u);}
|
||||
static METAL_FUNC float convert(threadgroup const bf16 & u) { return float(u);}
|
||||
static METAL_FUNC float convert(thread const bf16 & u) { return float(u);}
|
||||
};
|
||||
template<> struct convertor<bf16, float> {
|
||||
// fpext bfloat %_ to float
|
||||
static METAL_FUNC bf16 convert(device const float & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(threadgroup const float & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(thread const float & u) { return bf16(u); }
|
||||
};
|
||||
template<> struct convertor<float2, bf16_2> {
|
||||
// tail call fast <2 x float> @air.convert.f.v2f32.f.v2bf16(<2 x bfloat> %_)
|
||||
static METAL_FUNC float2 convert(device const bf16_2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(threadgroup const bf16_2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(thread const bf16_2 & u) { return float2(u); }
|
||||
};
|
||||
template<> struct convertor<bf16_2, float2> {
|
||||
// tail call fast <2 x bfloat> @air.convert.f.v2bf16.f.v2f32(<2 x float> %_)
|
||||
static METAL_FUNC bf16_2 convert(device const float2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(threadgroup const float2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(thread const float2 & u) { return bf16_2(u); }
|
||||
};
|
||||
|
||||
template<> struct convertor<float, half> {
|
||||
// fptrunc float %_ to half
|
||||
static METAL_FUNC float convert(device const half & u) { return float(u); }
|
||||
static METAL_FUNC float convert(threadgroup const half & u) { return float(u); }
|
||||
static METAL_FUNC float convert(thread const half & u) { return float(u); }
|
||||
};
|
||||
template<> struct convertor<half, float> {
|
||||
//fpext half %_ to float
|
||||
static METAL_FUNC half convert(device const float & u) { return half(u); }
|
||||
static METAL_FUNC half convert(threadgroup const float & u) { return half(u); }
|
||||
static METAL_FUNC half convert(thread const float & u) { return half(u); }
|
||||
};
|
||||
template<> struct convertor<float2, half2> {
|
||||
// tail call fast <2 x float> @air.convert.f.v2f32.f.v2f16(<2 x half> %_)
|
||||
static METAL_FUNC float2 convert(device const half2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(threadgroup const half2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(thread const half2 & u) { return float2(u); }
|
||||
};
|
||||
template<> struct convertor<half2, float2> {
|
||||
// tail call fast <2 x half> @air.convert.f.v2f16.f.v2f32(<2 x float> %_)
|
||||
static METAL_FUNC half2 convert(device const float2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(threadgroup const float2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(thread const float2 & u) { return half2(u); }
|
||||
};
|
||||
template<> struct convertor<bf16, half> {
|
||||
static METAL_FUNC bf16 convert(device const half & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(threadgroup const half & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(thread const half & u) { return bf16(u); }
|
||||
};
|
||||
template<> struct convertor<half, bf16> {
|
||||
static METAL_FUNC half convert(device const bf16 & u) { return half(u); }
|
||||
static METAL_FUNC half convert(threadgroup const bf16 & u) { return half(u); }
|
||||
static METAL_FUNC half convert(thread const bf16 & u) { return half(u); }
|
||||
};
|
||||
template<> struct convertor<bf16_2, half2> {
|
||||
// tail call fast <2 x bfloat> @air.convert.f.v2bf16.f.v2f16(<2 x half> %_)
|
||||
static METAL_FUNC bf16_2 convert(device const half2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(threadgroup const half2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(thread const half2 & u) { return bf16_2(u); }
|
||||
};
|
||||
template<> struct convertor<half2, bf16_2> {
|
||||
// tail call fast <2 x half> @air.convert.f.v2f16.f.v2bf16(<2 x bfloat> %_)
|
||||
static METAL_FUNC half2 convert(device const bf16_2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(threadgroup const bf16_2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(thread const bf16_2 & u) { return half2(u); }
|
||||
};
|
||||
|
||||
|
||||
|
||||
} // base_types
|
||||
|
||||
} // mittens
|
||||
@@ -0,0 +1,10 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief A collection of common resources on which Thundermittens depends.
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
#include "base_types.metal"
|
||||
#include "base_ops.metal"
|
||||
#include "utils.metal"
|
||||
@@ -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