diff --git a/.github/ci_config.yml b/.github/ci_config.yml index 68e2e00fa..eeeb383e7 100644 --- a/.github/ci_config.yml +++ b/.github/ci_config.yml @@ -149,6 +149,41 @@ platforms: - name: test run: pytest tests/ --devices cambricon -n 4 -v --tb=short --junitxml=/workspace/results/test-results.xml + hygon: + runner_label: Hygon + execution_mode: agent_local + image: + dockerfile: images/hygon/ + build_args: + BASE_IMAGE: image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.4.1-ubuntu22.04-dtk25.04.1-py3.10 + APT_MIRROR: http://archive.ubuntu.com/ubuntu + PIP_INDEX_URL: https://pypi.org/simple + docker_args: + - "--privileged" + - "--network=host" + - "--ipc=host" + - "--device=/dev/kfd" + - "--device=/dev/mkfd" + - "--device=/dev/dri" + - "--group-add=video" + volumes: + - /opt/hyhal:/opt/hyhal:ro + setup: pip install .[dev] --no-build-isolation + jobs: + gpu: + type: unittest + resources: + ngpus: 1 + gpu_style: none + memory: 32GB + shm_size: 64g + timeout: 3600 + queue_timeout: 600 + junit_path: test-results.xml + stages: + - name: test + run: pytest tests/ --devices hygon -n 4 -v --tb=short --junitxml=/workspace/results/test-results.xml + ascend: runner_label: Ascend execution_mode: agent_local diff --git a/CMakeLists.txt b/CMakeLists.txt index 739bea06e..3ac4bd400 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -15,6 +15,7 @@ set(PYBIND11_ENABLE_EXTRAS ON) option(WITH_CPU "Enable CPU backend" OFF) option(WITH_NVIDIA "Enable CUDA backend" OFF) option(WITH_ILUVATAR "Enable Iluvatar GPU backend" OFF) +option(WITH_HYGON "Enable Hygon GPU backend" OFF) option(WITH_METAX "Enable MetaX backend" OFF) option(WITH_CAMBRICON "Enable Cambricon backend" OFF) option(WITH_MOORE "Enable Moore backend" OFF) @@ -33,6 +34,31 @@ option(AUTO_DETECT_DEVICES "Automatically detect available devices" OFF) option(AUTO_DETECT_BACKENDS "Automatically detect available backends" OFF) option(GENERATE_PYTHON_BINDINGS "Generate Python bindings" OFF) +set(_DEFAULT_HYGON_DTK_ROOT "/opt/dtk") + +function(_infiniops_find_hygon_cuda_root out_var dtk_root) + set(_candidates + "${dtk_root}/cuda" + "${dtk_root}/cuda/cuda" + ) + + file(GLOB _versioned_cuda_dirs LIST_DIRECTORIES true "${dtk_root}/cuda/cuda-*") + if(_versioned_cuda_dirs) + list(SORT _versioned_cuda_dirs) + list(REVERSE _versioned_cuda_dirs) + list(APPEND _candidates ${_versioned_cuda_dirs}) + endif() + + foreach(_candidate IN LISTS _candidates) + if(EXISTS "${_candidate}/bin/nvcc") + set(${out_var} "${_candidate}" PARENT_SCOPE) + return() + endif() + endforeach() + + set(${out_var} "" PARENT_SCOPE) +endfunction() + if(AUTO_DETECT_DEVICES) message(STATUS "Auto-detecting available devices...") @@ -52,6 +78,24 @@ if(AUTO_DETECT_DEVICES) message(STATUS "Auto-detected Iluvatar environment.") endif() + set(_hygon_detected FALSE) + if(DEFINED ENV{DTK_ROOT} AND NOT "$ENV{DTK_ROOT}" STREQUAL "") + _infiniops_find_hygon_cuda_root(_HYGON_CUDA_DETECT_ROOT "$ENV{DTK_ROOT}") + if(_HYGON_CUDA_DETECT_ROOT) + set(_hygon_detected TRUE) + endif() + else() + _infiniops_find_hygon_cuda_root(_HYGON_CUDA_DETECT_ROOT "${_DEFAULT_HYGON_DTK_ROOT}") + if(_HYGON_CUDA_DETECT_ROOT) + set(_hygon_detected TRUE) + endif() + endif() + + if(_hygon_detected) + set(WITH_HYGON ON) + message(STATUS "Auto-detected Hygon environment.") + endif() + if(DEFINED ENV{MACA_PATH}) set(WITH_METAX ON) message(STATUS "Auto-detected MetaX environment from MACA_PATH") @@ -176,6 +220,17 @@ if(WITH_TORCH) OUTPUT_STRIP_TRAILING_WHITESPACE ) + execute_process( + COMMAND ${Python_EXECUTABLE} -c "import pathlib, torch; p = pathlib.Path(torch.__file__).resolve().parent.parent / 'torch.libs'; print(str(p) if p.exists() else '')" + OUTPUT_VARIABLE _torch_private_lib_dir + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + + set(TORCH_RUNTIME_DIRS ${_torch_lib_dirs}) + if(_torch_private_lib_dir) + list(APPEND TORCH_RUNTIME_DIRS ${_torch_private_lib_dir}) + endif() + find_library(TORCH_LIB torch HINTS ${_torch_lib_dirs} REQUIRED) find_library(TORCH_CPU_LIB torch_cpu HINTS ${_torch_lib_dirs} REQUIRED) find_library(C10_LIB c10 HINTS ${_torch_lib_dirs} REQUIRED) @@ -225,14 +280,14 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src) # Only one CUDA-like GPU backend can be enabled at a time. set(_gpu_backend_count 0) -foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_METAX WITH_MOORE WITH_ASCEND) +foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_METAX WITH_MOORE WITH_ASCEND) if(${_gpu_backend}) math(EXPR _gpu_backend_count "${_gpu_backend_count} + 1") endif() endforeach() if(_gpu_backend_count GREATER 1) - message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_METAX`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.") + message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_METAX`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.") endif() if(WITH_NVIDIA) @@ -261,6 +316,70 @@ if(WITH_ILUVATAR) find_package(CUDAToolkit REQUIRED) endif() +if(WITH_HYGON) + add_compile_definitions(WITH_HYGON=1) + set(DTK_ROOT $ENV{DTK_ROOT}) + if(NOT DTK_ROOT) + set(DTK_ROOT "${_DEFAULT_HYGON_DTK_ROOT}") + endif() + if(NOT EXISTS "${DTK_ROOT}") + message(FATAL_ERROR "`WITH_HYGON` is `ON` but `DTK_ROOT` (`${DTK_ROOT}`) does not exist.") + endif() + + set(_HYGON_ARCH_DEFAULT "gfx906") + if(DEFINED ENV{HYGON_ARCH} AND NOT "$ENV{HYGON_ARCH}" STREQUAL "") + set(_HYGON_ARCH_DEFAULT "$ENV{HYGON_ARCH}") + else() + find_program(HYGON_ROCMINFO_EXECUTABLE NAMES rocminfo HINTS "${DTK_ROOT}/bin") + if(HYGON_ROCMINFO_EXECUTABLE) + execute_process( + COMMAND ${HYGON_ROCMINFO_EXECUTABLE} + OUTPUT_VARIABLE _HYGON_ROCMINFO_OUTPUT + ERROR_QUIET + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + string(REGEX MATCH "gfx[0-9]+" _HYGON_ARCH_AUTO "${_HYGON_ROCMINFO_OUTPUT}") + if(_HYGON_ARCH_AUTO) + set(_HYGON_ARCH_DEFAULT "${_HYGON_ARCH_AUTO}") + endif() + endif() + endif() + + set(HYGON_ARCH "${_HYGON_ARCH_DEFAULT}" CACHE STRING "Hygon GPU architecture") + _infiniops_find_hygon_cuda_root(HYGON_CUDA_ROOT "${DTK_ROOT}") + + if(NOT HYGON_CUDA_ROOT) + message(FATAL_ERROR "`WITH_HYGON` is `ON` but no DTK `nvcc` was found under `${DTK_ROOT}`. Checked `${DTK_ROOT}/cuda/bin/nvcc`, `${DTK_ROOT}/cuda/cuda/bin/nvcc`, and `${DTK_ROOT}/cuda/cuda-*/bin/nvcc`.") + endif() + + set(CMAKE_CUDA_COMPILER "${HYGON_CUDA_ROOT}/bin/nvcc" CACHE FILEPATH "Hygon CUDA compiler (DTK nvcc)") + set(CUDAToolkit_ROOT "${HYGON_CUDA_ROOT}" CACHE PATH "Hygon CUDA toolkit root") + set(CMAKE_CUDA_ARCHITECTURES OFF CACHE STRING "Disable default CUDA arch flags for Hygon" FORCE) + set(CMAKE_CUDA_FLAGS "-std=c++17 -fPIC -arch=${HYGON_ARCH} -Wno-return-type -Wno-error=unused-private-field" CACHE STRING "Hygon CUDA flags") + set(CMAKE_CUDA_SEPARABLE_COMPILATION OFF CACHE BOOL "Disable RDC for Hygon") + + # DTK's nvcc wrapper derives its toolkit root from `CUDA_PATH`. + set(ENV{CUDA_PATH} "${HYGON_CUDA_ROOT}") + set(ENV{CUDA_HOME} "${HYGON_CUDA_ROOT}") + + # DTK's nvcc wrapper may invoke `nvcc` by name during compiler checks. + set(ENV{PATH} "${HYGON_CUDA_ROOT}/bin:$ENV{PATH}") + + # The actual Ninja build runs in fresh processes. Keep a launcher command + # for CUDA-backed Python bindings that need the DTK wrapper environment. + set(_HYGON_RULE_LAUNCH_ENV + "${CMAKE_COMMAND} -E env CUDA_PATH=${HYGON_CUDA_ROOT} CUDA_HOME=${HYGON_CUDA_ROOT} PATH=${HYGON_CUDA_ROOT}/bin:$ENV{PATH}") + + include_directories("${DTK_ROOT}/include") + include_directories("${HYGON_CUDA_ROOT}/include") + link_directories("${DTK_ROOT}/lib") + link_directories("${HYGON_CUDA_ROOT}/lib64") + + message(STATUS "Hygon: CUDA compiler ${CMAKE_CUDA_COMPILER}, arch ${HYGON_ARCH}, DTK root ${DTK_ROOT}") + enable_language(CUDA) + find_package(CUDAToolkit REQUIRED) +endif() + if(WITH_METAX) add_compile_definitions(WITH_METAX=1) @@ -344,7 +463,7 @@ if(WITH_ASCEND) endif() # If all other platforms are not enabled, CPU is enabled by default. -if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_METAX AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND) +if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_METAX AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND) add_compile_definitions(WITH_CPU=1) endif() @@ -352,6 +471,10 @@ if(WITH_TORCH OR WITH_METAX OR WITH_MOORE) set(PYBIND11_ENABLE_EXTRAS OFF) endif() +if(WITH_HYGON AND NOT EXISTS "${DTK_ROOT}/llvm/lib/LLVMgold.so") + set(PYBIND11_ENABLE_EXTRAS OFF) +endif() + add_subdirectory(src) if(NOT GENERATE_PYTHON_BINDINGS) diff --git a/README.md b/README.md index 6b9fc6f6e..44a8c71f1 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # InfiniOps -InfiniOps is a high-performance, cross-platform operator library supporting multiple backends: CPU, Nvidia, MetaX, Iluvatar, Moore, Cambricon, and more. +InfiniOps is a high-performance, cross-platform operator library supporting multiple backends: CPU, Nvidia, MetaX, Iluvatar, Hygon, Moore, Cambricon, and more. ## Prerequisites @@ -31,12 +31,16 @@ pip install . -C cmake.define.WITH_CPU=ON -C cmake.define.WITH_NVIDIA=ON | `-DWITH_NVIDIA=[ON\|OFF]` | Compile the Nvidia implementation | OFF | | `-DWITH_METAX=[ON\|OFF]` | Compile the MetaX implementation | OFF | | `-DWITH_ILUVATAR=[ON\|OFF]` | Compile the Iluvatar implementation | OFF | +| `-DWITH_HYGON=[ON\|OFF]` | Compile the Hygon implementation | OFF | | `-DWITH_MOORE=[ON\|OFF]` | Compile the Moore implementation | OFF | | `-DWITH_CAMBRICON=[ON\|OFF]` | Compile the Cambricon implementation | OFF | +| `-DWITH_ASCEND=[ON\|OFF]` | Compile the Ascend implementation | OFF | | `-DAUTO_DETECT_DEVICES=[ON\|OFF]` | Auto-detect available platforms | ON | If no accelerator options are provided and auto-detection finds nothing, `WITH_CPU` is enabled by default. +For Hygon builds, set `DTK_ROOT` to the DTK installation root if it is not installed at `/opt/dtk`. You can override the default DCU arch with `-DHYGON_ARCH=` when configuring CMake. + ## Contributing See [CONTRIBUTING.md](CONTRIBUTING.md) for code style, commit conventions, PR workflow, development guide, and troubleshooting. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 68ebc1b5c..18f700900 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -9,8 +9,14 @@ foreach(source_file ${EXAMPLE_SOURCES}) target_link_libraries(${example_name} PRIVATE infiniops) target_include_directories(${example_name} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) - + get_filename_component(example_dir ${source_file} DIRECTORY) target_include_directories(${example_name} PRIVATE ${example_dir}) + + if(WITH_TORCH) + foreach(_torch_dir ${TORCH_RUNTIME_DIRS}) + target_link_options(${example_name} PRIVATE "LINKER:-rpath-link,${_torch_dir}") + endforeach() + endif() endforeach() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 42cf3f6fb..4361ba38f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -70,6 +70,34 @@ if(WITH_ILUVATAR) list(APPEND DEVICE_LIST "iluvatar") endif() +if(WITH_HYGON) + set(HYGON_PATTERNS + "native/cuda/*.cc" + "native/cuda/*.cpp" + "native/cuda/*.cu" + "native/cuda/hygon/*.cc" + "native/cuda/hygon/*.cpp" + "native/cuda/hygon/*.cu" + ) + + file(GLOB_RECURSE HYGON_SOURCES CONFIGURE_DEPENDS ${HYGON_PATTERNS}) + + enable_language(CUDA) + + target_compile_definitions(infiniops PUBLIC WITH_HYGON=1) + target_sources(infiniops PRIVATE ${HYGON_SOURCES}) + + find_package(CUDAToolkit REQUIRED) + target_link_libraries(infiniops PUBLIC CUDA::cudart CUDA::cublas) + + set_target_properties(infiniops PROPERTIES + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + ) + + list(APPEND DEVICE_LIST "hygon") +endif() + if(WITH_METAX) set(METAX_PATTERNS "native/cuda/*.cc" @@ -517,7 +545,7 @@ if(GENERATE_PYTHON_BINDINGS) endif() list(APPEND PYBIND11_COMPILE_SOURCES ${PYBIND11_DISPATCH_SOURCES}) - if(WITH_NVIDIA) + if(WITH_NVIDIA OR WITH_HYGON) set_source_files_properties(${PYBIND11_COMPILE_SOURCES} PROPERTIES LANGUAGE CUDA) elseif(WITH_ILUVATAR) set(_iluvatar_dispatch_include_flags @@ -607,6 +635,13 @@ if(GENERATE_PYTHON_BINDINGS) target_compile_options(ops PRIVATE "-x" "musa") endif() + if(WITH_HYGON) + set_target_properties(ops PROPERTIES + RULE_LAUNCH_COMPILE "${_HYGON_RULE_LAUNCH_ENV}" + RULE_LAUNCH_LINK "${_HYGON_RULE_LAUNCH_ENV}" + ) + endif() + target_include_directories(ops PRIVATE ${PROJECT_SOURCE_DIR}) target_link_libraries(ops PRIVATE infiniops) @@ -620,9 +655,13 @@ if(GENERATE_PYTHON_BINDINGS) target_link_libraries(ops PRIVATE -Wl,--whole-archive no_workspace_kernel -Wl,--no-whole-archive) endif() + set(_INFINIOPS_INSTALL_RPATH "$ORIGIN") + if(WITH_TORCH) + list(APPEND _INFINIOPS_INSTALL_RPATH ${TORCH_RUNTIME_DIRS}) + endif() - set_target_properties(infiniops PROPERTIES INSTALL_RPATH "$ORIGIN") - set_target_properties(ops PROPERTIES INSTALL_RPATH "$ORIGIN") + set_target_properties(infiniops PROPERTIES INSTALL_RPATH "${_INFINIOPS_INSTALL_RPATH}") + set_target_properties(ops PROPERTIES INSTALL_RPATH "${_INFINIOPS_INSTALL_RPATH}") install(TARGETS infiniops ops DESTINATION .) diff --git a/src/native/cuda/hygon/blas.h b/src/native/cuda/hygon/blas.h new file mode 100644 index 000000000..dc9248cbf --- /dev/null +++ b/src/native/cuda/hygon/blas.h @@ -0,0 +1,40 @@ +#ifndef INFINI_OPS_HYGON_BLAS_H_ +#define INFINI_OPS_HYGON_BLAS_H_ + +#include + +// clang-format off +#include "cublas_v2.h" +// clang-format on + +#include "data_type.h" +#include "native/cuda/blas.h" +#include "native/cuda/hygon/blas_utils.h" +#include "native/cuda/hygon/runtime_.h" + +namespace infini::ops { + +template <> +struct Blas : public Runtime { + using BlasHandle = cublasHandle_t; + + static constexpr auto BLAS_OP_N = CUBLAS_OP_N; + + static constexpr auto BLAS_OP_T = CUBLAS_OP_T; + + static constexpr auto BLAS_GEMM_DEFAULT = CUBLAS_GEMM_DEFAULT_TENSOR_OP; + + static constexpr auto BlasCreate = cublasCreate; + + static constexpr auto BlasSetStream = cublasSetStream; + + static constexpr auto BlasDestroy = cublasDestroy; + + static constexpr auto BlasGemmStridedBatchedEx = [](auto&&... args) { + return cublasGemmStridedBatchedEx(std::forward(args)...); + }; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/hygon/blas_utils.h b/src/native/cuda/hygon/blas_utils.h new file mode 100644 index 000000000..7c55c5564 --- /dev/null +++ b/src/native/cuda/hygon/blas_utils.h @@ -0,0 +1,29 @@ +#ifndef INFINI_OPS_HYGON_BLAS_UTILS_H_ +#define INFINI_OPS_HYGON_BLAS_UTILS_H_ + +// clang-format off +#include "cublas_v2.h" +// clang-format on + +#include "data_type.h" +#include "native/cuda/blas_utils.h" + +namespace infini::ops { + +template <> +struct BlasUtils { + static auto GetDataType(DataType dtype) { + if (dtype == DataType::kFloat16) return CUDA_R_16F; + if (dtype == DataType::kBFloat16) return CUDA_R_16BF; + return CUDA_R_32F; + } + + static auto GetComputeType(DataType dtype) { + (void)dtype; + return CUBLAS_COMPUTE_32F; + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/hygon/device_.h b/src/native/cuda/hygon/device_.h new file mode 100644 index 000000000..1df1246e7 --- /dev/null +++ b/src/native/cuda/hygon/device_.h @@ -0,0 +1,67 @@ +#ifndef INFINI_OPS_HYGON_DEVICE__H_ +#define INFINI_OPS_HYGON_DEVICE__H_ + +// clang-format off +#include +#include +// clang-format on + +#include "data_type.h" +#include "device.h" +#include "native/cuda/caster.cuh" + +namespace infini::ops { + +template <> +struct DeviceEnabled : std::true_type {}; + +// Some DTK toolchains expose the underlying bf16 structs but gate the +// nv_bfloat16 typedefs behind CUDA_NO_BFLOAT16. +using cuda_bfloat16 = __nv_bfloat16; + +using cuda_bfloat162 = __nv_bfloat162; + +namespace detail { + +template <> +struct ToFloat { + __host__ __device__ float operator()(half x) { return __half2float(x); } +}; + +template <> +struct ToFloat { + __host__ __device__ float operator()(__nv_bfloat16 x) { + return __bfloat162float(x); + } +}; + +template <> +struct FromFloat { + __host__ __device__ half operator()(float f) { return __float2half(f); } +}; + +template <> +struct FromFloat { + __host__ __device__ __nv_bfloat16 operator()(float f) { + return __float2bfloat16(f); + } +}; + +} // namespace detail + +template <> +struct TypeMap { + using type = half; +}; + +template <> +struct TypeMap { + using type = __nv_bfloat16; +}; + +template <> +struct Caster : CudaCasterImpl {}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/hygon/device_property.h b/src/native/cuda/hygon/device_property.h new file mode 100644 index 000000000..28e6bce7a --- /dev/null +++ b/src/native/cuda/hygon/device_property.h @@ -0,0 +1,42 @@ +#ifndef INFINI_OPS_HYGON_DEVICE_PROPERTY_H_ +#define INFINI_OPS_HYGON_DEVICE_PROPERTY_H_ + +#include + +#include +#include + +namespace infini::ops { + +class DevicePropertyCache { + public: + static const cudaDeviceProp& GetCurrentDeviceProps() { + int device_id = 0; + cudaGetDevice(&device_id); + return GetDeviceProps(device_id); + } + + static const cudaDeviceProp& GetDeviceProps(int device_id) { + static std::vector cache = []() { + int count = 0; + cudaGetDeviceCount(&count); + if (count == 0) return std::vector{}; + std::vector props(count); + for (int i = 0; i < count; ++i) { + cudaGetDeviceProperties(&props[i], i); + } + return props; + }(); + + assert(device_id >= 0 && device_id < static_cast(cache.size())); + return cache[device_id]; + } +}; + +inline int QueryMaxThreadsPerBlock() { + return DevicePropertyCache::GetCurrentDeviceProps().maxThreadsPerBlock; +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/hygon/runtime_.h b/src/native/cuda/hygon/runtime_.h new file mode 100644 index 000000000..da1f63e7d --- /dev/null +++ b/src/native/cuda/hygon/runtime_.h @@ -0,0 +1,44 @@ +#ifndef INFINI_OPS_HYGON_RUNTIME_H_ +#define INFINI_OPS_HYGON_RUNTIME_H_ + +#include + +// clang-format off +#include +// clang-format on + +#include "native/cuda/hygon/device_.h" +#include "native/cuda/hygon/runtime_utils.h" +#include "native/cuda/runtime_.h" + +namespace infini::ops { + +template <> +struct Runtime + : CudaRuntime> { + using Stream = cudaStream_t; + + static constexpr Device::Type kDeviceType = Device::Type::kHygon; + + static constexpr auto Malloc = [](auto&&... args) { + return cudaMalloc(std::forward(args)...); + }; + + static constexpr auto Memcpy = cudaMemcpy; + + static constexpr auto Free = [](auto&&... args) { + return cudaFree(std::forward(args)...); + }; + + static constexpr auto MemcpyHostToDevice = cudaMemcpyHostToDevice; + + static constexpr auto MemcpyDeviceToHost = cudaMemcpyDeviceToHost; + + static constexpr auto Memset = cudaMemset; +}; + +static_assert(Runtime::Validate()); + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/hygon/runtime_utils.h b/src/native/cuda/hygon/runtime_utils.h new file mode 100644 index 000000000..e5604ac81 --- /dev/null +++ b/src/native/cuda/hygon/runtime_utils.h @@ -0,0 +1,15 @@ +#ifndef INFINI_OPS_HYGON_RUNTIME_UTILS_H_ +#define INFINI_OPS_HYGON_RUNTIME_UTILS_H_ + +#include "native/cuda/hygon/device_property.h" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template <> +struct RuntimeUtils + : CudaRuntimeUtils {}; + +} // namespace infini::ops + +#endif diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 0f6332d86..b76dc1a2f 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -4,6 +4,8 @@ #include #include +#include + #include "tensor.h" #include "torch/device_.h" @@ -38,7 +40,37 @@ inline Device::Type DeviceTypeFromString(const std::string& name) { return it->second; } - return Device::TypeFromString(name); + std::vector supported_names; + + for (const auto& [torch_name, device_type] : kTorchNameToTypes) { + const auto internal_name = std::string{Device::StringFromType(device_type)}; + + if (name == internal_name) { + return device_type; + } + + supported_names.push_back(torch_name); + supported_names.push_back(internal_name); + } + + std::sort(supported_names.begin(), supported_names.end()); + supported_names.erase( + std::unique(supported_names.begin(), supported_names.end()), + supported_names.end()); + + std::string message = "Unsupported device type `" + name + + "` for this InfiniOps build. Supported device names: "; + + for (std::size_t i = 0; i < supported_names.size(); ++i) { + if (i != 0) { + message += ", "; + } + message += supported_names[i]; + } + + message += ". Rebuild InfiniOps with the matching backend enabled."; + + throw py::value_error(message); } // Returns `nullopt` rather than aborting when the name does not resolve. diff --git a/tests/conftest.py b/tests/conftest.py index b38c7f574..301a7da3f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,7 @@ def pytest_addoption(parser): "--devices", nargs="+", default=None, - help="Device(s) to test on (e.g., `--devices ascend cpu`). Accepts platform names (`nvidia`, `metax`, `iluvatar`, `moore`, `cambricon`, `ascend`) or PyTorch device types (`cuda`, `mlu`, `musa`, `npu`). Defaults to all available devices.", + help="Device(s) to test on (e.g., `--devices ascend cpu`). Accepts platform names (`nvidia`, `metax`, `iluvatar`, `hygon`, `moore`, `cambricon`, `ascend`) or PyTorch device types (`cuda`, `mlu`, `musa`, `npu`). Defaults to all available devices.", ) @@ -138,6 +138,7 @@ def _set_random_seed(seed): "nvidia": "cuda", "metax": "cuda", "iluvatar": "cuda", + "hygon": "cuda", "moore": "musa", "cambricon": "mlu", "ascend": "npu",