Skip to content

Commit 2c9455a

Browse files
ZiminliZiminli
authored
feat: support unified cast and type-safe host math fallbacks for reduce-type MPI operations (#38)
* feat: add the generic `Caster` in `src/caster.h` and apply it in the reduce-related MPI implementations - add the generic `Caster` in `src/caster.h` - apply the generic caster in the reduce-related MPI implementations (i.e., `all_reduce`, `reduce`, and `reduce_scatter`) - update the bridge file generation logic to also include the platform-specific caster files and refactor it to search for a flexible combination * feat: add CPU-specialized `HardwareCastImpl` * feat: add NVIDIA-specialized `HardwareCastImpl` * feat: add Iluvatar-specialized `HardwareCastImpl` * feat: add MetaX-specialized `HardwareCastImpl` * feat: add Moore-specialized `HardwareCastImpl` * feat: add Cambricon-specialized `HardwareCastImpl` and update `data_type_.h` to properly represent fp16 and bf16 * feat: add `SupportsOp` trait for checking operation supportiveness at compile-time, provide convenience wrappers, and apply them to the reduce-related mpi implementations * fix: explicitly disable math operations on `half` and `__nv_bfloat16` on Iluvatar since they are not really supported on the host side. * fix: make `AUTO_DETECT_DEVICES` for NVIDIA to be specific by explicitly looking for the number 0 GPU card * feat: support compiling device code on Iluvatar * fix: explicitly disable math operations on `half` and `__mt_bfloat16` on MThreads since they are not really supported on the host side. * style: clang-format `src/cambricon/caster.mlu`
1 parent 86e29fc commit 2c9455a

17 files changed

Lines changed: 551 additions & 36 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ if(AUTO_DETECT_DEVICES)
3838
# NVIDIA
3939
set(NVIDIA_FOUND FALSE)
4040

41-
file(GLOB NVIDIA_DEV_FILES "/dev/nvidia*")
41+
file(GLOB NVIDIA_DEV_FILES "/dev/nvidia0")
4242

4343
if(NVIDIA_DEV_FILES)
4444
set(NVIDIA_FOUND TRUE)

‎examples/CMakeLists.txt‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,22 @@ foreach(source_file ${EXAMPLE_SOURCES})
88
target_link_libraries(${example_name} PRIVATE infiniccl)
99

1010
# Add runtime and backend dependencies for direct runtime/backend usage.
11-
if(WITH_NVIDIA OR WITH_ILUVATAR)
11+
if(WITH_NVIDIA)
1212
target_link_libraries(${example_name} PRIVATE CUDA::cudart)
1313
endif()
1414

15+
if(WITH_ILUVATAR)
16+
set_source_files_properties(${source_file} PROPERTIES LANGUAGE CXX)
17+
set_target_properties(${example_name} PROPERTIES
18+
RULE_LAUNCH_COMPILE "${ILUVATAR_CUDA_COMPILER} "
19+
)
20+
target_compile_options(${example_name} PRIVATE ${ILUVATAR_CUDA_FLAGS} "-Wno-unused-command-line-argument")
21+
target_link_libraries(${example_name} PRIVATE CUDA::cudart CUDA::cuda_driver)
22+
endif()
23+
1524
if(WITH_METAX)
1625
target_link_libraries(${example_name} PRIVATE ${MACA_RUNTIME_LIB})
26+
target_compile_options(${example_name} PRIVATE "-x" "maca")
1727
endif()
1828

1929
if(WITH_MOORE)

‎scripts/gen_bridge.py‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,12 @@
1919
"""
2020

2121
# Hardware traits to look for in device directories.
22-
DEVICE_TRAIT_HEADERS = ["device_.h", "runtime_.h", "data_type_.h"]
22+
DEVICE_TRAIT_FILES = {
23+
"caster_": ["h", "cuh"],
24+
"device_": ["h"],
25+
"runtime_": ["h", "cuh"],
26+
"data_type_": ["h", "cuh"],
27+
}
2328

2429
# Map logical backend names (from CMake) to their internal source paths.
2530
BACKEND_PATH_MAP = {"ompi": "ompi/impl", "mpich": "ompi/impl", "nccl": "nvidia/nccl"}
@@ -97,11 +102,14 @@ def generate(project_root, output_dir, devices, backends):
97102
device_included = False
98103
manifest_lines.append(f"\n// --- DEVICE: {dev.upper()} ---")
99104

100-
for trait in DEVICE_TRAIT_HEADERS:
101-
rel_path = f"{dev}/{trait}"
102-
if os.path.exists(os.path.join(src_dir, rel_path)):
103-
manifest_lines.append(f'#include "{rel_path}"')
104-
device_included = True
105+
for trait, extensions in DEVICE_TRAIT_FILES.items():
106+
for ext in extensions:
107+
rel_path = f"{dev}/{trait}.{ext}"
108+
109+
if os.path.exists(os.path.join(src_dir, rel_path)):
110+
manifest_lines.append(f'#include "{rel_path}"')
111+
device_included = True
112+
break
105113

106114
if device_included:
107115
found_devices.append(f"Device::Type::k{dev.capitalize()}")

‎src/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ if(WITH_ILUVATAR)
7373
set_source_files_properties(${ILUVATAR_SOURCES} PROPERTIES
7474
COMPILE_OPTIONS "${ILUVATAR_CUDA_FLAGS}"
7575
)
76+
set(CMAKE_CXX_COMPILER "${ILUVATAR_CUDA_COMPILER}" CACHE FILEPATH "Iluvatar Compiler" FORCE)
77+
target_compile_options(infiniccl PRIVATE ${ILUVATAR_CUDA_FLAGS})
7678

7779
target_link_libraries(infiniccl PRIVATE CUDA::cudart CUDA::cuda_driver)
7880
endif()

‎src/cambricon/caster.mlu‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
#include "caster_.h"
2+
3+
namespace infini::ccl {
4+
5+
float HardwareCastImpl<Device::Type::kCambricon, float, CambriconFP16>::Apply(
6+
CambriconFP16 x) {
7+
const __half& native_x = *reinterpret_cast<const __half*>(&x.bits);
8+
return __half2float(native_x);
9+
}
10+
11+
CambriconFP16 HardwareCastImpl<Device::Type::kCambricon, CambriconFP16,
12+
float>::Apply(float x) {
13+
__half native_res = __float2half__(x);
14+
CambriconFP16 out;
15+
out.bits = *reinterpret_cast<uint16_t*>(&native_res);
16+
return out;
17+
}
18+
19+
CambriconFP16
20+
HardwareCastImpl<Device::Type::kCambricon, CambriconFP16, int>::Apply(int x) {
21+
__half native_res = __int2half_rn__(x);
22+
CambriconFP16 out;
23+
out.bits = *reinterpret_cast<uint16_t*>(&native_res);
24+
return out;
25+
}
26+
27+
float HardwareCastImpl<Device::Type::kCambricon, float, CambriconBF16>::Apply(
28+
CambriconBF16 x) {
29+
const bfloat16_t& native_x = *reinterpret_cast<const bfloat16_t*>(&x.bits);
30+
return __bfloat162float__(native_x);
31+
}
32+
33+
CambriconBF16 HardwareCastImpl<Device::Type::kCambricon, CambriconBF16,
34+
float>::Apply(float x) {
35+
bfloat16_t native_res = __float2bfloat16__(x);
36+
CambriconBF16 out;
37+
out.bits = *reinterpret_cast<uint16_t*>(&native_res);
38+
return out;
39+
}
40+
41+
CambriconBF16
42+
HardwareCastImpl<Device::Type::kCambricon, CambriconBF16, int>::Apply(int x) {
43+
bfloat16_t native_res = __int2bfloat16_rn__(x);
44+
CambriconBF16 out;
45+
out.bits = *reinterpret_cast<uint16_t*>(&native_res);
46+
return out;
47+
}
48+
49+
} // namespace infini::ccl

‎src/cambricon/caster_.h‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
#ifndef INFINI_CCL_CAMBRICON_CASTER__H_
2+
#define INFINI_CCL_CAMBRICON_CASTER__H_
3+
4+
#include "caster.h"
5+
#include "data_type_.h"
6+
7+
namespace infini::ccl {
8+
9+
template <>
10+
struct HardwareCastImpl<Device::Type::kCambricon, float, CambriconFP16> {
11+
static float Apply(CambriconFP16 x);
12+
};
13+
14+
template <>
15+
struct HardwareCastImpl<Device::Type::kCambricon, CambriconFP16, float> {
16+
static CambriconFP16 Apply(float x);
17+
};
18+
19+
template <>
20+
struct HardwareCastImpl<Device::Type::kCambricon, CambriconFP16, int> {
21+
static CambriconFP16 Apply(int x);
22+
};
23+
24+
template <>
25+
struct HardwareCastImpl<Device::Type::kCambricon, float, CambriconBF16> {
26+
static float Apply(CambriconBF16 x);
27+
};
28+
29+
template <>
30+
struct HardwareCastImpl<Device::Type::kCambricon, CambriconBF16, float> {
31+
static CambriconBF16 Apply(float x);
32+
};
33+
34+
template <>
35+
struct HardwareCastImpl<Device::Type::kCambricon, CambriconBF16, int> {
36+
static CambriconBF16 Apply(int x);
37+
};
38+
39+
} // namespace infini::ccl
40+
41+
#endif // INFINI_CCL_CAMBRICON_CASTER__H_

‎src/cambricon/data_type_.h‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,28 @@
11
#ifndef INFINI_CCL_CAMBRICON_DATA_TYPE__H_
22
#define INFINI_CCL_CAMBRICON_DATA_TYPE__H_
33

4-
#if defined(__mlu__) || defined(__BANG__)
5-
#include "bang_bf16.h"
6-
#else
7-
// If compiling on the host side using standard GCC, mock the storage types
8-
// to prevent header poisoning and avoid `bang_fp16.h` dependency altogether.
9-
#ifndef HAS_CAMBRICON_HOST_HALF_MOCKS
10-
#define HAS_CAMBRICON_HOST_HALF_MOCKS
4+
#include <cstdint>
115

12-
typedef uint16_t half;
13-
typedef uint16_t bfloat16_t;
14-
15-
#endif
16-
#endif
17-
18-
#include "cambricon/device_.h"
196
#include "data_type_impl.h"
207

218
namespace infini::ccl {
229

10+
// Unique tag types to force completely distinct template signatures.
11+
struct CambriconFP16 {
12+
uint16_t bits;
13+
};
14+
struct CambriconBF16 {
15+
uint16_t bits;
16+
};
17+
2318
template <>
2419
struct TypeMap<Device::Type::kCambricon, DataType::kFloat16> {
25-
using type = half;
20+
using type = CambriconFP16;
2621
};
2722

2823
template <>
2924
struct TypeMap<Device::Type::kCambricon, DataType::kBFloat16> {
30-
using type = bfloat16_t;
25+
using type = CambriconBF16;
3126
};
3227

3328
} // namespace infini::ccl

‎src/caster.h‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
#ifndef INFINI_CCL_CASTER_H_
2+
#define INFINI_CCL_CASTER_H_
3+
4+
#include <type_traits>
5+
#include <typeinfo>
6+
#include <utility>
7+
8+
#include "device.h"
9+
10+
namespace infini::ccl {
11+
12+
// Check if a type is complete (i.e. fully defined) at compile-time.
13+
template <typename T, typename = void>
14+
struct IsComplete : std::false_type {};
15+
16+
template <typename T>
17+
struct IsComplete<T, std::void_t<decltype(sizeof(T))>> : std::true_type {};
18+
19+
template <Device::Type kDev, typename Dst, typename Src>
20+
struct HardwareCastImpl;
21+
22+
template <Device::Type kDev, typename Dst, typename Src>
23+
inline constexpr bool HasHardwareCastV =
24+
IsComplete<HardwareCastImpl<kDev, Dst, Src>>::value;
25+
26+
// Intermediate type used for fallback conversions, usually `float`.
27+
template <Device::Type kDev>
28+
struct CastBridge {
29+
using Type = float;
30+
};
31+
32+
template <Device::Type kDev>
33+
struct Caster {
34+
template <typename Dst, typename Src>
35+
static constexpr Dst Cast(Src&& x) {
36+
using PureSrc = std::remove_cv_t<std::remove_reference_t<Src>>;
37+
using PureDst = std::remove_cv_t<std::remove_reference_t<Dst>>;
38+
using Bridge = typename CastBridge<kDev>::Type;
39+
40+
if constexpr (std::is_same_v<PureSrc, PureDst>) {
41+
return std::forward<Src>(x);
42+
} else if constexpr (HasHardwareCastV<kDev, PureDst, PureSrc>) {
43+
return HardwareCastImpl<kDev, PureDst, PureSrc>::Apply(
44+
std::forward<Src>(x));
45+
} else if constexpr (!std::is_same_v<PureSrc, Bridge> &&
46+
!std::is_same_v<PureDst, Bridge> &&
47+
HasHardwareCastV<kDev, Bridge, PureSrc> &&
48+
(HasHardwareCastV<kDev, PureDst, Bridge> ||
49+
std::is_arithmetic_v<PureDst>)) {
50+
Bridge tmp = Cast<Bridge>(std::forward<Src>(x));
51+
return Cast<PureDst>(tmp);
52+
} else if constexpr (!std::is_same_v<PureDst, Bridge> &&
53+
std::is_arithmetic_v<PureSrc> &&
54+
HasHardwareCastV<kDev, PureDst, Bridge>) {
55+
Bridge tmp = static_cast<Bridge>(std::forward<Src>(x));
56+
return Cast<PureDst>(tmp);
57+
} else if constexpr (std::is_arithmetic_v<PureSrc> &&
58+
std::is_arithmetic_v<PureDst>) {
59+
return static_cast<PureDst>(std::forward<Src>(x));
60+
} else {
61+
static_assert(HasHardwareCastV<kDev, PureDst, PureSrc>,
62+
"no cast path available. "
63+
"Need to provide `HardwareCastImpl` specialization.");
64+
}
65+
}
66+
};
67+
68+
// Convenience wrapper for casting between two types.
69+
// Otherwise, need to write `Caster<kDev>::template Cast<Target>(val)` every
70+
// time.
71+
template <Device::Type kDev, typename Target, typename Source>
72+
inline Target CastTo(Source&& val) {
73+
return Caster<kDev>::template Cast<Target>(std::forward<Source>(val));
74+
}
75+
76+
template <Device::Type kDev, typename Source>
77+
inline float ToFloat(Source&& val) {
78+
return Caster<kDev>::template Cast<float>(std::forward<Source>(val));
79+
}
80+
81+
} // namespace infini::ccl
82+
83+
#endif // INFINI_CCL_CASTER_H_

‎src/cpu/caster_.h‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
#ifndef INFINI_CCL_CPU_CASTER__H_
2+
#define INFINI_CCL_CPU_CASTER__H_
3+
4+
#include "caster.h"
5+
#include "data_type_.h"
6+
7+
namespace infini::ccl {
8+
9+
template <>
10+
struct HardwareCastImpl<Device::Type::kCpu, float, Float16> {
11+
static float Apply(Float16 x) { return x.ToFloat(); }
12+
};
13+
14+
template <>
15+
struct HardwareCastImpl<Device::Type::kCpu, Float16, float> {
16+
static Float16 Apply(float x) { return Float16::FromFloat(x); }
17+
};
18+
19+
template <>
20+
struct HardwareCastImpl<Device::Type::kCpu, float, BFloat16> {
21+
static float Apply(BFloat16 x) { return x.ToFloat(); }
22+
};
23+
24+
template <>
25+
struct HardwareCastImpl<Device::Type::kCpu, BFloat16, float> {
26+
static BFloat16 Apply(float x) { return BFloat16::FromFloat(x); }
27+
};
28+
29+
} // namespace infini::ccl
30+
31+
#endif // INFINI_CCL_CPU_CASTER__H_

‎src/iluvatar/caster_.cuh‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#ifndef INFINI_CCL_ILUVATAR_CASTER__CUH_
2+
#define INFINI_CCL_ILUVATAR_CASTER__CUH_
3+
4+
#include "caster.h"
5+
#include "data_type_.h"
6+
#include "traits.h"
7+
8+
namespace infini::ccl {
9+
10+
// Explicitly tell the trait system that host-side math on these types is
11+
// unworkable on Iluvatar.
12+
template <typename S, typename Op>
13+
struct SupportsOp<half, S, Op, void> : std::false_type {};
14+
15+
template <typename S, typename Op>
16+
struct SupportsOp<__nv_bfloat16, S, Op, void> : std::false_type {};
17+
18+
template <>
19+
struct HardwareCastImpl<Device::Type::kIluvatar, float, half> {
20+
__host__ __device__ static float Apply(half x) { return __half2float(x); }
21+
};
22+
23+
template <>
24+
struct HardwareCastImpl<Device::Type::kIluvatar, float, __nv_bfloat16> {
25+
__host__ __device__ static float Apply(__nv_bfloat16 x) {
26+
return __bfloat162float(x);
27+
}
28+
};
29+
30+
template <>
31+
struct HardwareCastImpl<Device::Type::kIluvatar, half, float> {
32+
__host__ __device__ static half Apply(float x) { return __float2half(x); }
33+
};
34+
35+
template <>
36+
struct HardwareCastImpl<Device::Type::kIluvatar, __nv_bfloat16, float> {
37+
__host__ __device__ static __nv_bfloat16 Apply(float x) {
38+
return __float2bfloat16(x);
39+
}
40+
};
41+
42+
template <>
43+
struct HardwareCastImpl<Device::Type::kIluvatar, __nv_bfloat16, double> {
44+
__host__ __device__ static __nv_bfloat16 Apply(double x) {
45+
return __double2bfloat16(x);
46+
}
47+
};
48+
49+
} // namespace infini::ccl
50+
51+
#endif // INFINI_CCL_ILUVATAR_CASTER__CUH_

0 commit comments

Comments
 (0)