MLX
Loading...
Searching...
No Matches
Classes | Macros | Typedefs | Functions
utils.h File Reference
#include <metal_math>
#include "mlx/backend/metal/kernels/bf16.h"
#include "mlx/backend/metal/kernels/complex.h"
#include "mlx/backend/metal/kernels/defines.h"

Go to the source code of this file.

Classes

struct  Limits< U >
 
struct  Limits< uint8_t >
 
struct  Limits< uint16_t >
 
struct  Limits< uint32_t >
 
struct  Limits< uint64_t >
 
struct  Limits< int8_t >
 
struct  Limits< int16_t >
 
struct  Limits< int32_t >
 
struct  Limits< int64_t >
 
struct  Limits< half >
 
struct  Limits< float >
 
struct  Limits< bfloat16_t >
 
struct  Limits< bool >
 
struct  Limits< complex64_t >
 
struct  looped_elem_to_loc< dim, offset_t >
 
struct  looped_elem_to_loc< 1, offset_t >
 
struct  looped_elem_to_loc< 0, offset_t >
 

Macros

#define instantiate_default_limit(type)
 
#define instantiate_float_limit(type)
 
#define MLX_MTL_PRAGMA_UNROLL   _Pragma("clang loop unroll(full)")
 

Typedefs

typedef half float16_t
 

Functions

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc (uint elem, device const int *shape, device const stride_t *strides, int ndim)
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc (uint elem, constant const int *shape, constant const stride_t *strides, int ndim)
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc (stride_t elem, device const int *shape, device const stride_t *strides, int ndim)
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc (stride_t elem, constant const int *shape, constant const stride_t *strides, int ndim)
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc (uint3 elem, constant const int *shape, constant const stride_t *strides, int ndim)
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc_1 (uint elem, constant const stride_t &stride)
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc_2 (uint2 elem, constant const stride_t strides[2])
 
template<typename stride_t >
METAL_FUNC stride_t elem_to_loc_3 (uint3 elem, constant const stride_t strides[3])
 
template<int NDIM>
METAL_FUNC size_t elem_to_loc_nd (uint elem, device const int *shape, device const size_t *strides)
 
template<int NDIM>
METAL_FUNC size_t elem_to_loc_nd (uint3 elem, constant const int shape[NDIM], constant const size_t strides[NDIM])
 
template<int NDIM>
METAL_FUNC int64_t elem_to_loc_nd (uint elem, constant const int shape[NDIM], constant const int64_t strides[NDIM])
 
template<int NDIM>
METAL_FUNC int64_t elem_to_loc_nd (uint3 elem, constant const int shape[NDIM], constant const int64_t strides[NDIM])
 
METAL_FUNC uint2 elem_to_loc_2_nd (uint3 elem, constant const int *shape, constant const size_t *a_strides, constant const size_t *b_strides, int ndim)
 
METAL_FUNC uint3 elem_to_loc_3_nd (uint3 elem, constant const int *shape, constant const size_t *a_strides, constant const size_t *b_strides, constant const size_t *c_strides, int ndim)
 
template<int NDIM>
METAL_FUNC uint2 elem_to_loc_2_nd (uint3 elem, constant const int shape[NDIM], constant const size_t a_strides[NDIM], constant const size_t b_strides[NDIM])
 
template<int NDIM>
METAL_FUNC uint3 elem_to_loc_3_nd (uint3 elem, constant const int shape[NDIM], constant const size_t a_strides[NDIM], constant const size_t b_strides[NDIM], constant const size_t c_strides[NDIM])
 
template<typename T , typename U >
T ceildiv (T N, U M)
 Compute ceil((float)N/(float)M)
 
float log1p (float x)
 
bfloat16_t log1p (bfloat16_t x)
 
uint64_t simd_shuffle_down (uint64_t data, uint16_t delta)
 
int64_t simd_shuffle_down (int64_t data, uint16_t delta)
 
bool simd_shuffle_down (bool data, uint16_t delta)
 
complex64_t simd_shuffle_down (complex64_t data, uint16_t delta)
 

Macro Definition Documentation

◆ instantiate_default_limit

#define instantiate_default_limit ( type)
Value:
template <> \
struct Limits<type> { \
static constexpr constant type max = metal::numeric_limits<type>::max(); \
static constexpr constant type min = metal::numeric_limits<type>::min(); \
static constexpr constant type finite_max = \
metal::numeric_limits<type>::max(); \
static constexpr constant type finite_min = \
metal::numeric_limits<type>::min(); \
};
Limits
Definition utils.h:17
Limits::max
static const constant U max
Definition utils.h:18
Limits::finite_max
static const constant U finite_max
Definition utils.h:20
Limits::min
static const constant U min
Definition utils.h:19
Limits::finite_min
static const constant U finite_min
Definition utils.h:21

◆ instantiate_float_limit

#define instantiate_float_limit ( type)
Value:
template <> \
struct Limits<type> { \
static constexpr constant type max = \
metal::numeric_limits<type>::infinity(); \
static constexpr constant type min = \
-metal::numeric_limits<type>::infinity(); \
static constexpr constant type finite_max = \
metal::numeric_limits<type>::max(); \
static constexpr constant type finite_min = \
-metal::numeric_limits<type>::max(); \
};

◆ MLX_MTL_PRAGMA_UNROLL

#define MLX_MTL_PRAGMA_UNROLL   _Pragma("clang loop unroll(full)")

Typedef Documentation

◆ float16_t

typedef half float16_t

Function Documentation

◆ ceildiv()

template<typename T , typename U >
T ceildiv ( T N,
U M )
inline

Compute ceil((float)N/(float)M)

◆ elem_to_loc() [1/5]

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc ( stride_t elem,
constant const int * shape,
constant const stride_t * strides,
int ndim )

◆ elem_to_loc() [2/5]

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc ( stride_t elem,
device const int * shape,
device const stride_t * strides,
int ndim )

◆ elem_to_loc() [3/5]

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc ( uint elem,
constant const int * shape,
constant const stride_t * strides,
int ndim )

◆ elem_to_loc() [4/5]

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc ( uint elem,
device const int * shape,
device const stride_t * strides,
int ndim )

◆ elem_to_loc() [5/5]

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc ( uint3 elem,
constant const int * shape,
constant const stride_t * strides,
int ndim )

◆ elem_to_loc_1()

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc_1 ( uint elem,
constant const stride_t & stride )

◆ elem_to_loc_2()

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc_2 ( uint2 elem,
constant const stride_t strides[2] )

◆ elem_to_loc_2_nd() [1/2]

METAL_FUNC uint2 elem_to_loc_2_nd ( uint3 elem,
constant const int * shape,
constant const size_t * a_strides,
constant const size_t * b_strides,
int ndim )

◆ elem_to_loc_2_nd() [2/2]

template<int NDIM>
METAL_FUNC uint2 elem_to_loc_2_nd ( uint3 elem,
constant const int shape[NDIM],
constant const size_t a_strides[NDIM],
constant const size_t b_strides[NDIM] )

◆ elem_to_loc_3()

template<typename stride_t >
METAL_FUNC stride_t elem_to_loc_3 ( uint3 elem,
constant const stride_t strides[3] )

◆ elem_to_loc_3_nd() [1/2]

METAL_FUNC uint3 elem_to_loc_3_nd ( uint3 elem,
constant const int * shape,
constant const size_t * a_strides,
constant const size_t * b_strides,
constant const size_t * c_strides,
int ndim )

◆ elem_to_loc_3_nd() [2/2]

template<int NDIM>
METAL_FUNC uint3 elem_to_loc_3_nd ( uint3 elem,
constant const int shape[NDIM],
constant const size_t a_strides[NDIM],
constant const size_t b_strides[NDIM],
constant const size_t c_strides[NDIM] )

◆ elem_to_loc_nd() [1/4]

template<int NDIM>
METAL_FUNC int64_t elem_to_loc_nd ( uint elem,
constant const int shape[NDIM],
constant const int64_t strides[NDIM] )

◆ elem_to_loc_nd() [2/4]

template<int NDIM>
METAL_FUNC size_t elem_to_loc_nd ( uint elem,
device const int * shape,
device const size_t * strides )

◆ elem_to_loc_nd() [3/4]

template<int NDIM>
METAL_FUNC int64_t elem_to_loc_nd ( uint3 elem,
constant const int shape[NDIM],
constant const int64_t strides[NDIM] )

◆ elem_to_loc_nd() [4/4]

template<int NDIM>
METAL_FUNC size_t elem_to_loc_nd ( uint3 elem,
constant const int shape[NDIM],
constant const size_t strides[NDIM] )

◆ log1p() [1/2]

bfloat16_t log1p ( bfloat16_t x)
inline

◆ log1p() [2/2]

float log1p ( float x)
inline

◆ simd_shuffle_down() [1/4]

bool simd_shuffle_down ( bool data,
uint16_t delta )
inline

◆ simd_shuffle_down() [2/4]

complex64_t simd_shuffle_down ( complex64_t data,
uint16_t delta )
inline

◆ simd_shuffle_down() [3/4]

int64_t simd_shuffle_down ( int64_t data,
uint16_t delta )
inline

◆ simd_shuffle_down() [4/4]

uint64_t simd_shuffle_down ( uint64_t data,
uint16_t delta )
inline
Generated by doxygen 1.10.0