mirror of
https://github.com/ml-explore/mlx.git
synced 2025-12-16 01:49:05 +08:00
MLX_SWITCH macros to templates (#2320)
This commit is contained in:
Angelos Katharopoulos
committed by
GitHub
GitHub
parent
33bf1a244b
commit
3d5e17e507
@@ -1,207 +1,106 @@
|
||||
// Copyright © 2025 Apple Inc.
|
||||
// Copyright © Meta Platforms, Inc. and affiliates.
|
||||
//
|
||||
// This source code is licensed under the BSD-style license found in
|
||||
// https://github.com/pytorch/executorch/blob/main/LICENSE
|
||||
//
|
||||
// Forked from
|
||||
// https://github.com/pytorch/executorch/blob/main/runtime/core/exec_aten/util/scalar_type_util.h
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "mlx/dtype.h"
|
||||
#include <sstream>
|
||||
|
||||
#include <fmt/format.h>
|
||||
#include "mlx/dtype.h"
|
||||
#include "mlx/utils.h"
|
||||
|
||||
namespace mlx::core {
|
||||
|
||||
// Return string representation of dtype.
|
||||
const char* dtype_to_string(Dtype arg);
|
||||
|
||||
// Macros that iterate across different subsets of Dtypes.
|
||||
//
|
||||
// For all of these macros, the final `_` parameter is the name of another macro
|
||||
// that takes two parameters: the name of a C type, and the name of the
|
||||
// corresponding Dtype enumerator.
|
||||
//
|
||||
// Note that these macros should use fully-qualified namespaces (starting with
|
||||
// `::`) to ensure that they can be called safely in any arbitrary namespace.
|
||||
#define MLX_FORALL_INT_TYPES(_) \
|
||||
_(uint8_t, uint8) \
|
||||
_(uint16_t, uint16) \
|
||||
_(uint32_t, uint32) \
|
||||
_(uint64_t, uint64) \
|
||||
_(int8_t, int8) \
|
||||
_(int16_t, int16) \
|
||||
_(int32_t, int32) \
|
||||
_(int64_t, int64)
|
||||
#define MLX_INTERNAL_DTYPE_SWITCH_CASE(DTYPE, TYPE) \
|
||||
case DTYPE: \
|
||||
f(type_identity<TYPE>{}); \
|
||||
break
|
||||
|
||||
#define MLX_FORALL_FLOAT_TYPES(_) \
|
||||
_(float16_t, float16) \
|
||||
_(float, float32) \
|
||||
_(double, float64) \
|
||||
_(bfloat16_t, bfloat16)
|
||||
#define MLX_INTERNAL_DTYPE_SWITCH_INTS() \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(int8, int8_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(int16, int16_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(int32, int32_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(int64, int64_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(uint8, uint8_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(uint16, uint16_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(uint32, uint32_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(uint64, uint64_t)
|
||||
|
||||
// Calls the provided macro on every Dtype, providing the C type and the
|
||||
// Dtype name to each call.
|
||||
//
|
||||
// @param _ A macro that takes two parameters: the name of a C type, and the
|
||||
// name of the corresponding Dtype enumerator.
|
||||
#define MLX_FORALL_DTYPES(_) \
|
||||
MLX_FORALL_INT_TYPES(_) \
|
||||
MLX_FORALL_FLOAT_TYPES(_) \
|
||||
_(bool, bool_) \
|
||||
_(complex64_t, complex64)
|
||||
#define MLX_INTERNAL_DTYPE_SWITCH_FLOATS() \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(float16, float16_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(bfloat16, bfloat16_t); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(float32, float); \
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(float64, double)
|
||||
|
||||
// Maps Dtypes to C++ types.
|
||||
template <Dtype::Val N>
|
||||
struct DtypeToCppType;
|
||||
|
||||
#define SPECIALIZE_DtypeToCppType(CPP_TYPE, DTYPE) \
|
||||
template <> \
|
||||
struct DtypeToCppType<Dtype::Val::DTYPE> { \
|
||||
using type = CPP_TYPE; \
|
||||
};
|
||||
|
||||
MLX_FORALL_DTYPES(SPECIALIZE_DtypeToCppType)
|
||||
|
||||
#undef SPECIALIZE_DtypeToCppType
|
||||
|
||||
// Maps C++ types to Dtypes.
|
||||
// This already exists in C++20 but in C++20 we can also just use templated
|
||||
// lambdas which will make this so much nicer.
|
||||
template <typename T>
|
||||
struct CppTypeToDtype;
|
||||
struct type_identity {
|
||||
using type = T;
|
||||
};
|
||||
|
||||
#define SPECIALIZE_CppTypeToDtype(CPP_TYPE, DTYPE) \
|
||||
template <> \
|
||||
struct CppTypeToDtype<CPP_TYPE> \
|
||||
: std::integral_constant<Dtype::Val, Dtype::Val::DTYPE> {};
|
||||
#define MLX_GET_TYPE(x) typename decltype(x)::type
|
||||
#define MLX_GET_VALUE(x) decltype(x)::value
|
||||
|
||||
MLX_FORALL_DTYPES(SPECIALIZE_CppTypeToDtype)
|
||||
|
||||
#undef SPECIALIZE_CppTypeToDtype
|
||||
|
||||
// Helper macros for switch case macros (see below)
|
||||
//
|
||||
// These macros are not meant to be used directly. They provide an easy way to
|
||||
// generate a switch statement that can handle subsets of Dtypes supported.
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE(enum_type, CTYPE_ALIAS, ...) \
|
||||
case enum_type: { \
|
||||
using CTYPE_ALIAS = ::mlx::core::DtypeToCppType<enum_type>::type; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
template <typename F>
|
||||
void dispatch_all_types(Dtype dt, F&& f) {
|
||||
switch (dt) {
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(bool_, bool);
|
||||
MLX_INTERNAL_DTYPE_SWITCH_INTS();
|
||||
MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(complex64, complex64_t);
|
||||
}
|
||||
}
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CHECKED(TYPE, NAME, ...) \
|
||||
switch (TYPE) { \
|
||||
__VA_ARGS__ \
|
||||
default: \
|
||||
throw std::invalid_argument(fmt::format( \
|
||||
"Unhandled dtype %s for %s", dtype_to_string(TYPE), NAME)); \
|
||||
template <typename F>
|
||||
void dispatch_int_types(Dtype dt, std::string_view tag, F&& f) {
|
||||
switch (dt) {
|
||||
MLX_INTERNAL_DTYPE_SWITCH_INTS();
|
||||
default:
|
||||
std::ostringstream msg;
|
||||
msg << tag << " Only integer types supported but " << dt
|
||||
<< " was provided";
|
||||
throw std::invalid_argument(msg.str());
|
||||
}
|
||||
}
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE_INT_TYPES(CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::uint8, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::uint16, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::uint32, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::uint64, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::int8, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::int16, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::int32, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::int64, CTYPE_ALIAS, __VA_ARGS__)
|
||||
template <typename F>
|
||||
void dispatch_float_types(Dtype dt, std::string_view tag, F&& f) {
|
||||
switch (dt) {
|
||||
MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
|
||||
default:
|
||||
std::ostringstream msg;
|
||||
msg << tag << " Only float types supported but " << dt << " was provided";
|
||||
throw std::invalid_argument(msg.str());
|
||||
}
|
||||
}
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE_FLOAT_TYPES(CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::float16, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::float32, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::float64, CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::bfloat16, CTYPE_ALIAS, __VA_ARGS__)
|
||||
template <typename F>
|
||||
void dispatch_int_float_types(Dtype dt, std::string_view tag, F&& f) {
|
||||
switch (dt) {
|
||||
MLX_INTERNAL_DTYPE_SWITCH_INTS();
|
||||
MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
|
||||
default:
|
||||
std::ostringstream msg;
|
||||
msg << tag << " Only integer and float types supported but " << dt
|
||||
<< " was provided";
|
||||
throw std::invalid_argument(msg.str());
|
||||
}
|
||||
}
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE_INT_FLOAT_TYPES(CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CASE_INT_TYPES(CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE_FLOAT_TYPES(CTYPE_ALIAS, __VA_ARGS__)
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE_REAL_TYPES(CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CASE_INT_FLOAT_TYPES(CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::bool_, CTYPE_ALIAS, __VA_ARGS__)
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE_COMPLEX_TYPES(CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CASE( \
|
||||
::mlx::core::Dtype::Val::complex64, CTYPE_ALIAS, __VA_ARGS__)
|
||||
|
||||
#define MLX_INTERNAL_SWITCH_CASE_ALL_TYPES(CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CASE_REAL_TYPES(CTYPE_ALIAS, __VA_ARGS__) \
|
||||
MLX_INTERNAL_SWITCH_CASE_COMPLEX_TYPES(CTYPE_ALIAS, __VA_ARGS__)
|
||||
|
||||
// Switch case macros
|
||||
//
|
||||
// These macros provide an easy way to generate switch statements that apply a
|
||||
// common lambda function to subsets of Dtypes supported by MLX.
|
||||
// The lambda function can type specialize to the ctype associated with the
|
||||
// Dtype being handled through an alias passed as the CTYPE_ALIAS argument.
|
||||
//
|
||||
// Arguments:
|
||||
// - ADDITIONAL: Additional Dtype case to add
|
||||
// - TYPE: The Dtype to handle through the switch statement
|
||||
// - NAME: A name for this operation which will be used in error messages
|
||||
// - CTYPE_ALIAS: A typedef for the ctype associated with the Dtype.
|
||||
// - ...: A statement to be applied to each Dtype case
|
||||
//
|
||||
// An example usage is:
|
||||
//
|
||||
// MLX_SWITCH_ALL_TYPES(input.dtype(), CTYPE, {
|
||||
// output.data<CTYPE>[0] = input.data<CTYPE>[0];
|
||||
// });
|
||||
//
|
||||
// Note that these can be nested as well:
|
||||
//
|
||||
// MLX_SWITCH_ALL_TYPES(input.dtype(), CTYPE_IN, {
|
||||
// MLX_SWITCH_ALL_TYPES(output.dtype(), CTYPE_OUT, {
|
||||
// output.data<CTYPE_OUT>[0] = input.data<CTYPE_IN>[0];
|
||||
// });
|
||||
// });
|
||||
//
|
||||
// These macros are adapted from Dispatch.h in the ATen library. The primary
|
||||
// difference is that the CTYPE_ALIAS argument is exposed to users, which is
|
||||
// used to alias the ctype associated with the Dtype that is being handled.
|
||||
|
||||
#define MLX_SWITCH_ALL_TYPES(TYPE, CTYPE_ALIAS, ...) \
|
||||
switch (TYPE) { MLX_INTERNAL_SWITCH_CASE_ALL_TYPES(CTYPE_ALIAS, __VA_ARGS__) }
|
||||
|
||||
#define MLX_SWITCH_INT_TYPES_CHECKED(TYPE, NAME, CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CHECKED( \
|
||||
TYPE, \
|
||||
NAME, \
|
||||
MLX_INTERNAL_SWITCH_CASE_INT_TYPES(CTYPE_ALIAS, __VA_ARGS__))
|
||||
|
||||
#define MLX_SWITCH_FLOAT_TYPES_CHECKED(TYPE, NAME, CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CHECKED( \
|
||||
TYPE, \
|
||||
NAME, \
|
||||
MLX_INTERNAL_SWITCH_CASE_FLOAT_TYPES(CTYPE_ALIAS, __VA_ARGS__))
|
||||
|
||||
#define MLX_SWITCH_INT_FLOAT_TYPES_CHECKED(TYPE, NAME, CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CHECKED( \
|
||||
TYPE, \
|
||||
NAME, \
|
||||
MLX_INTERNAL_SWITCH_CASE_INT_FLOAT_TYPES(CTYPE_ALIAS, __VA_ARGS__))
|
||||
|
||||
#define MLX_SWITCH_REAL_TYPES_CHECKED(TYPE, NAME, CTYPE_ALIAS, ...) \
|
||||
MLX_INTERNAL_SWITCH_CHECKED( \
|
||||
TYPE, \
|
||||
NAME, \
|
||||
MLX_INTERNAL_SWITCH_CASE_REAL_TYPES(CTYPE_ALIAS, __VA_ARGS__))
|
||||
template <typename F>
|
||||
void dispatch_real_types(Dtype dt, std::string_view tag, F&& f) {
|
||||
switch (dt) {
|
||||
MLX_INTERNAL_DTYPE_SWITCH_CASE(bool_, bool);
|
||||
MLX_INTERNAL_DTYPE_SWITCH_INTS();
|
||||
MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
|
||||
default:
|
||||
std::ostringstream msg;
|
||||
msg << tag << " Only real numbers supported but " << dt
|
||||
<< " was provided";
|
||||
throw std::invalid_argument(msg.str());
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mlx::core
|
||||
|
||||
Reference in New Issue
Block a user