diff --git a/include/infini/ops.h b/include/infini/ops.h index db17bd33..5da15661 100644 --- a/include/infini/ops.h +++ b/include/infini/ops.h @@ -1,8 +1,120 @@ #ifndef INFINI_OPS_H_ #define INFINI_OPS_H_ +#include +#include + #ifdef __cplusplus #include + +extern "C" { +#endif + +#if defined(_WIN32) +#if defined(INFINI_OPS_BUILD_SHARED) +#define INFINI_OPS_API __declspec(dllexport) +#elif defined(INFINI_OPS_USE_SHARED) +#define INFINI_OPS_API __declspec(dllimport) +#else +#define INFINI_OPS_API +#endif +#else +#if defined(INFINI_OPS_BUILD_SHARED) +#define INFINI_OPS_API __attribute__((visibility("default"))) +#else +#define INFINI_OPS_API +#endif +#endif + +typedef enum InfiniOpsStatus { + INFINI_OPS_STATUS_SUCCESS = 0, + INFINI_OPS_STATUS_INVALID_ARGUMENT = 1, + INFINI_OPS_STATUS_NOT_SUPPORTED = 2, + INFINI_OPS_STATUS_OUT_OF_MEMORY = 3, + INFINI_OPS_STATUS_INTERNAL_ERROR = 4, +} InfiniOpsStatus; + +typedef enum InfiniOpsDataType { + INFINI_OPS_DATA_TYPE_INVALID = 0, + INFINI_OPS_DATA_TYPE_FLOAT16 = 1, + INFINI_OPS_DATA_TYPE_FLOAT32 = 2, + INFINI_OPS_DATA_TYPE_FLOAT64 = 3, + INFINI_OPS_DATA_TYPE_BFLOAT16 = 4, + INFINI_OPS_DATA_TYPE_INT8 = 5, + INFINI_OPS_DATA_TYPE_INT16 = 6, + INFINI_OPS_DATA_TYPE_INT32 = 7, + INFINI_OPS_DATA_TYPE_INT64 = 8, + INFINI_OPS_DATA_TYPE_UINT8 = 9, + INFINI_OPS_DATA_TYPE_UINT16 = 10, + INFINI_OPS_DATA_TYPE_UINT32 = 11, + INFINI_OPS_DATA_TYPE_UINT64 = 12, +} InfiniOpsDataType; + +typedef enum InfiniOpsDeviceType { + INFINI_OPS_DEVICE_TYPE_INVALID = 0, + INFINI_OPS_DEVICE_TYPE_CPU = 1, + INFINI_OPS_DEVICE_TYPE_NVIDIA = 2, + INFINI_OPS_DEVICE_TYPE_CAMBRICON = 3, + INFINI_OPS_DEVICE_TYPE_ASCEND = 4, + INFINI_OPS_DEVICE_TYPE_METAX = 5, + INFINI_OPS_DEVICE_TYPE_MOORE = 6, + INFINI_OPS_DEVICE_TYPE_ILUVATAR = 7, + INFINI_OPS_DEVICE_TYPE_KUNLUN = 8, + INFINI_OPS_DEVICE_TYPE_HYGON = 9, + INFINI_OPS_DEVICE_TYPE_QY = 10, +} InfiniOpsDeviceType; + +typedef struct InfiniOpsTensor { + size_t structure_size; + void *data; + size_t byte_size; + InfiniOpsDataType data_type; + InfiniOpsDeviceType device_type; + int32_t rank; + const int64_t *shape; + const int64_t *stride; + uint64_t reserved[8]; +} InfiniOpsTensor; + +typedef struct InfiniOpsStreamPrivate *InfiniOpsStream; +typedef struct InfiniOpsHandlePrivate *InfiniOpsHandle; +typedef struct InfiniOpsConfigPrivate *InfiniOpsConfig; + +typedef struct InfiniOpsHandleAttributes { + size_t structure_size; + InfiniOpsStream stream; + void *workspace; + size_t workspace_byte_size; + uint64_t reserved[8]; +} InfiniOpsHandleAttributes; + +typedef struct InfiniOpsConfigAttributes { + size_t structure_size; + size_t implementation_index; + uint64_t reserved[8]; +} InfiniOpsConfigAttributes; + +INFINI_OPS_API InfiniOpsStatus infiniOpsGetLastError(char *buffer, + size_t capacity, + size_t *required_size); + +INFINI_OPS_API InfiniOpsStatus infiniOpsCreateHandle( + const InfiniOpsHandleAttributes *attributes, InfiniOpsHandle *handle); + +INFINI_OPS_API InfiniOpsStatus infiniOpsDestroyHandle(InfiniOpsHandle handle); + +INFINI_OPS_API InfiniOpsStatus infiniOpsCreateConfig( + const InfiniOpsConfigAttributes *attributes, InfiniOpsConfig *config); + +INFINI_OPS_API InfiniOpsStatus infiniOpsDestroyConfig(InfiniOpsConfig config); + +INFINI_OPS_API InfiniOpsStatus infiniOpsAdd( + InfiniOpsHandle handle, InfiniOpsConfig config, + const InfiniOpsTensor *input, const InfiniOpsTensor *other, + InfiniOpsTensor *output); + +#ifdef __cplusplus +} #endif #endif // INFINI_OPS_H_ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6086e072..b74fc9a2 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -2,6 +2,30 @@ add_library(infiniops SHARED) file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc") target_sources(infiniops PRIVATE ${BASE_SRCS}) +target_sources(infiniops PRIVATE infini/ops.cc) +target_compile_definitions(infiniops PRIVATE INFINI_OPS_BUILD_SHARED=1) + +target_include_directories(infiniops + PUBLIC + $ + $ + $ + $ +) + +set_target_properties(infiniops PROPERTIES + VERSION 1.0.0 + SOVERSION 1 +) + +if(UNIX AND NOT APPLE AND CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(infiniops PRIVATE + "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/infini/ops.map" + ) + set_target_properties(infiniops PROPERTIES + LINK_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/infini/ops.map" + ) +endif() set(DEVICE_LIST "") @@ -446,14 +470,6 @@ if(WITH_TORCH) endif() endif() -target_include_directories(infiniops - PUBLIC - $ - $ - $ - $ -) - if(GENERATE_CPP_OPERATOR_API OR GENERATE_PYTHON_BINDINGS) find_package(Python COMPONENTS Interpreter REQUIRED) # Always regenerate wrappers so the generated functional API and pybind11 diff --git a/src/infini/ops.cc b/src/infini/ops.cc new file mode 100644 index 00000000..5fdd74aa --- /dev/null +++ b/src/infini/ops.cc @@ -0,0 +1,494 @@ +#include "infini/ops.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +struct InfiniOpsHandlePrivate { + infini::ops::Handle handle; +}; + +struct InfiniOpsConfigPrivate { + infini::ops::Config config; +}; + +namespace { + +thread_local std::string last_error; + +void SetLastError(const char *message) { last_error = message; } + +InfiniOpsStatus InvalidArgument(const char *message) { + SetLastError(message); + return INFINI_OPS_STATUS_INVALID_ARGUMENT; +} + +bool IsValidDataType(InfiniOpsDataType data_type) { + switch (data_type) { + case INFINI_OPS_DATA_TYPE_FLOAT16: + case INFINI_OPS_DATA_TYPE_BFLOAT16: + case INFINI_OPS_DATA_TYPE_FLOAT32: + case INFINI_OPS_DATA_TYPE_FLOAT64: + case INFINI_OPS_DATA_TYPE_INT8: + case INFINI_OPS_DATA_TYPE_INT16: + case INFINI_OPS_DATA_TYPE_INT32: + case INFINI_OPS_DATA_TYPE_INT64: + case INFINI_OPS_DATA_TYPE_UINT8: + case INFINI_OPS_DATA_TYPE_UINT16: + case INFINI_OPS_DATA_TYPE_UINT32: + case INFINI_OPS_DATA_TYPE_UINT64: + return true; + case INFINI_OPS_DATA_TYPE_INVALID: + return false; + } + return false; +} + +bool IsValidDeviceType(InfiniOpsDeviceType device_type) { + switch (device_type) { + case INFINI_OPS_DEVICE_TYPE_CPU: + case INFINI_OPS_DEVICE_TYPE_NVIDIA: + case INFINI_OPS_DEVICE_TYPE_CAMBRICON: + case INFINI_OPS_DEVICE_TYPE_ASCEND: + case INFINI_OPS_DEVICE_TYPE_METAX: + case INFINI_OPS_DEVICE_TYPE_MOORE: + case INFINI_OPS_DEVICE_TYPE_ILUVATAR: + case INFINI_OPS_DEVICE_TYPE_KUNLUN: + case INFINI_OPS_DEVICE_TYPE_HYGON: + case INFINI_OPS_DEVICE_TYPE_QY: + return true; + case INFINI_OPS_DEVICE_TYPE_INVALID: + return false; + } + return false; +} + +bool ConvertDataType(InfiniOpsDataType input, infini::ops::DataType *output) { + assert(output != nullptr); + switch (input) { + case INFINI_OPS_DATA_TYPE_FLOAT16: + *output = infini::ops::DataType::kFloat16; + return true; + case INFINI_OPS_DATA_TYPE_FLOAT32: + *output = infini::ops::DataType::kFloat32; + return true; + case INFINI_OPS_DATA_TYPE_FLOAT64: + *output = infini::ops::DataType::kFloat64; + return true; + case INFINI_OPS_DATA_TYPE_BFLOAT16: + *output = infini::ops::DataType::kBFloat16; + return true; + case INFINI_OPS_DATA_TYPE_INT8: + *output = infini::ops::DataType::kInt8; + return true; + case INFINI_OPS_DATA_TYPE_INT16: + *output = infini::ops::DataType::kInt16; + return true; + case INFINI_OPS_DATA_TYPE_INT32: + *output = infini::ops::DataType::kInt32; + return true; + case INFINI_OPS_DATA_TYPE_INT64: + *output = infini::ops::DataType::kInt64; + return true; + case INFINI_OPS_DATA_TYPE_UINT8: + *output = infini::ops::DataType::kUInt8; + return true; + case INFINI_OPS_DATA_TYPE_UINT16: + *output = infini::ops::DataType::kUInt16; + return true; + case INFINI_OPS_DATA_TYPE_UINT32: + *output = infini::ops::DataType::kUInt32; + return true; + case INFINI_OPS_DATA_TYPE_UINT64: + *output = infini::ops::DataType::kUInt64; + return true; + case INFINI_OPS_DATA_TYPE_INVALID: + return false; + } + return false; +} + +bool ConvertDeviceType(InfiniOpsDeviceType input, + infini::ops::Device::Type *output) { + assert(output != nullptr); + switch (input) { + case INFINI_OPS_DEVICE_TYPE_CPU: + *output = infini::ops::Device::Type::kCpu; + return true; + case INFINI_OPS_DEVICE_TYPE_NVIDIA: + *output = infini::ops::Device::Type::kNvidia; + return true; + case INFINI_OPS_DEVICE_TYPE_CAMBRICON: + *output = infini::ops::Device::Type::kCambricon; + return true; + case INFINI_OPS_DEVICE_TYPE_ASCEND: + *output = infini::ops::Device::Type::kAscend; + return true; + case INFINI_OPS_DEVICE_TYPE_METAX: + *output = infini::ops::Device::Type::kMetax; + return true; + case INFINI_OPS_DEVICE_TYPE_MOORE: + *output = infini::ops::Device::Type::kMoore; + return true; + case INFINI_OPS_DEVICE_TYPE_ILUVATAR: + *output = infini::ops::Device::Type::kIluvatar; + return true; + case INFINI_OPS_DEVICE_TYPE_KUNLUN: + *output = infini::ops::Device::Type::kKunlun; + return true; + case INFINI_OPS_DEVICE_TYPE_HYGON: + *output = infini::ops::Device::Type::kHygon; + return true; + case INFINI_OPS_DEVICE_TYPE_QY: + *output = infini::ops::Device::Type::kQy; + return true; + case INFINI_OPS_DEVICE_TYPE_INVALID: + return false; + } + return false; +} + +bool DataTypeSize(InfiniOpsDataType data_type, size_t *size) { + assert(size != nullptr); + switch (data_type) { + case INFINI_OPS_DATA_TYPE_FLOAT16: + case INFINI_OPS_DATA_TYPE_BFLOAT16: + case INFINI_OPS_DATA_TYPE_INT16: + case INFINI_OPS_DATA_TYPE_UINT16: + *size = 2; + return true; + case INFINI_OPS_DATA_TYPE_FLOAT32: + case INFINI_OPS_DATA_TYPE_INT32: + case INFINI_OPS_DATA_TYPE_UINT32: + *size = 4; + return true; + case INFINI_OPS_DATA_TYPE_FLOAT64: + case INFINI_OPS_DATA_TYPE_INT64: + case INFINI_OPS_DATA_TYPE_UINT64: + *size = 8; + return true; + case INFINI_OPS_DATA_TYPE_INT8: + case INFINI_OPS_DATA_TYPE_UINT8: + *size = 1; + return true; + case INFINI_OPS_DATA_TYPE_INVALID: + return false; + } + return false; +} + +bool CheckedMultiply(size_t left, size_t right, size_t *result) { + assert(result != nullptr); + if (right != 0 && left > SIZE_MAX / right) { + return false; + } + *result = left * right; + return true; +} + +bool ComputeRequiredByteSize(const InfiniOpsTensor &tensor, + size_t *required_byte_size) { + assert(required_byte_size != nullptr); + + size_t element_size = 0; + if (!DataTypeSize(tensor.data_type, &element_size)) { + return false; + } + + if (tensor.rank == 0) { + *required_byte_size = element_size; + return true; + } + + size_t element_count = 1; + if (tensor.stride == nullptr) { + for (int32_t i = 0; i < tensor.rank; ++i) { + if (!CheckedMultiply(element_count, static_cast(tensor.shape[i]), + &element_count)) { + return false; + } + } + return CheckedMultiply(element_count, element_size, required_byte_size); + } + + size_t max_offset = 0; + for (int32_t i = 0; i < tensor.rank; ++i) { + if (tensor.shape[i] == 0) { + *required_byte_size = 0; + return true; + } + if (tensor.stride[i] < 0) { + return false; + } + size_t dimension_extent = 0; + if (!CheckedMultiply(static_cast(tensor.shape[i] - 1), + static_cast(tensor.stride[i]), + &dimension_extent)) { + return false; + } + if (SIZE_MAX - max_offset < dimension_extent) { + return false; + } + max_offset += dimension_extent; + } + + if (max_offset == SIZE_MAX) { + return false; + } + return CheckedMultiply(max_offset + 1, element_size, required_byte_size); +} + +std::vector ConvertShape( + const InfiniOpsTensor &tensor) { + std::vector result; + result.reserve(tensor.rank); + for (int32_t i = 0; i < tensor.rank; ++i) { + result.push_back(static_cast(tensor.shape[i])); + } + return result; +} + +std::vector ConvertStrides( + const InfiniOpsTensor &tensor) { + std::vector result; + result.reserve(tensor.rank); + if (tensor.stride == nullptr) { + return result; + } + for (int32_t i = 0; i < tensor.rank; ++i) { + result.push_back(static_cast(tensor.stride[i])); + } + return result; +} + +infini::ops::Tensor ToInternalTensor(const InfiniOpsTensor &tensor) { + infini::ops::DataType data_type; + const bool data_type_valid = ConvertDataType(tensor.data_type, &data_type); + assert(data_type_valid); + + infini::ops::Device::Type device_type; + const bool device_type_valid = + ConvertDeviceType(tensor.device_type, &device_type); + assert(device_type_valid); + + const auto shape = ConvertShape(tensor); + const infini::ops::Device device(device_type); + if (tensor.stride == nullptr) { + return infini::ops::Tensor(tensor.data, shape, data_type, device); + } + return infini::ops::Tensor(tensor.data, shape, data_type, device, + ConvertStrides(tensor)); +} + +InfiniOpsStatus ValidateTensor(const char *name, + const InfiniOpsTensor *tensor) { + if (tensor == nullptr) { + SetLastError(name); + last_error += " tensor must not be null"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (tensor->structure_size < offsetof(InfiniOpsTensor, reserved)) { + SetLastError(name); + last_error += " tensor structure size is invalid"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (!IsValidDataType(tensor->data_type)) { + SetLastError(name); + last_error += " tensor data type is invalid"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (!IsValidDeviceType(tensor->device_type)) { + SetLastError(name); + last_error += " tensor device type is invalid"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (tensor->rank < 0) { + SetLastError(name); + last_error += " tensor rank must not be negative"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (tensor->rank > 0 && tensor->shape == nullptr) { + SetLastError(name); + last_error += " tensor shape must not be null for non-scalar"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + for (int32_t i = 0; i < tensor->rank; ++i) { + if (tensor->shape[i] < 0) { + SetLastError(name); + last_error += " tensor shape must not contain negative values"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + } + + size_t required_byte_size = 0; + if (!ComputeRequiredByteSize(*tensor, &required_byte_size)) { + SetLastError(name); + last_error += " tensor byte size is invalid"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (required_byte_size > 0 && tensor->data == nullptr) { + SetLastError(name); + last_error += " tensor data must not be null"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (tensor->byte_size < required_byte_size) { + SetLastError(name); + last_error += " tensor byte size is smaller than shape requires"; + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + return INFINI_OPS_STATUS_SUCCESS; +} + +} // namespace + +extern "C" { + +InfiniOpsStatus infiniOpsGetLastError(char *buffer, size_t capacity, + size_t *required_size) { + const size_t required = last_error.size() + 1; + if (required_size != nullptr) { + *required_size = required; + } + if (buffer == nullptr || capacity == 0) { + return last_error.empty() ? INFINI_OPS_STATUS_SUCCESS + : INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + if (capacity < required) { + if (capacity > 0) { + buffer[0] = '\0'; + } + return INFINI_OPS_STATUS_INVALID_ARGUMENT; + } + std::memcpy(buffer, last_error.c_str(), required); + return INFINI_OPS_STATUS_SUCCESS; +} + +InfiniOpsStatus infiniOpsCreateHandle( + const InfiniOpsHandleAttributes *attributes, InfiniOpsHandle *handle) { + try { + if (handle == nullptr) { + return InvalidArgument("handle output must not be null"); + } + *handle = nullptr; + if (attributes == nullptr) { + return InvalidArgument("handle attributes must not be null"); + } + if (attributes->structure_size < + offsetof(InfiniOpsHandleAttributes, reserved)) { + return InvalidArgument("handle attributes size is invalid"); + } + if (attributes->workspace_byte_size > 0 && + attributes->workspace == nullptr) { + return InvalidArgument("handle workspace must not be null"); + } + + InfiniOpsHandlePrivate *created = new InfiniOpsHandlePrivate; + created->handle.set_stream(reinterpret_cast(attributes->stream)); + created->handle.set_workspace(attributes->workspace); + created->handle.set_workspace_size_in_bytes( + attributes->workspace_byte_size); + *handle = created; + SetLastError(""); + return INFINI_OPS_STATUS_SUCCESS; + } catch (const std::bad_alloc &) { + SetLastError("out of memory while creating handle"); + return INFINI_OPS_STATUS_OUT_OF_MEMORY; + } catch (const std::exception &error) { + SetLastError(error.what()); + return INFINI_OPS_STATUS_INTERNAL_ERROR; + } catch (...) { + SetLastError("unknown error while creating handle"); + return INFINI_OPS_STATUS_INTERNAL_ERROR; + } +} + +InfiniOpsStatus infiniOpsDestroyHandle(InfiniOpsHandle handle) { + delete handle; + SetLastError(""); + return INFINI_OPS_STATUS_SUCCESS; +} + +InfiniOpsStatus infiniOpsCreateConfig( + const InfiniOpsConfigAttributes *attributes, InfiniOpsConfig *config) { + try { + if (config == nullptr) { + return InvalidArgument("config output must not be null"); + } + *config = nullptr; + if (attributes == nullptr) { + return InvalidArgument("config attributes must not be null"); + } + if (attributes->structure_size < + offsetof(InfiniOpsConfigAttributes, reserved)) { + return InvalidArgument("config attributes size is invalid"); + } + + InfiniOpsConfigPrivate *created = new InfiniOpsConfigPrivate; + created->config.set_implementation_index(attributes->implementation_index); + *config = created; + SetLastError(""); + return INFINI_OPS_STATUS_SUCCESS; + } catch (const std::bad_alloc &) { + SetLastError("out of memory while creating config"); + return INFINI_OPS_STATUS_OUT_OF_MEMORY; + } catch (const std::exception &error) { + SetLastError(error.what()); + return INFINI_OPS_STATUS_INTERNAL_ERROR; + } catch (...) { + SetLastError("unknown error while creating config"); + return INFINI_OPS_STATUS_INTERNAL_ERROR; + } +} + +InfiniOpsStatus infiniOpsDestroyConfig(InfiniOpsConfig config) { + delete config; + SetLastError(""); + return INFINI_OPS_STATUS_SUCCESS; +} + +InfiniOpsStatus infiniOpsAdd(InfiniOpsHandle handle, InfiniOpsConfig config, + const InfiniOpsTensor *input, + const InfiniOpsTensor *other, + InfiniOpsTensor *output) { + try { + InfiniOpsStatus status = ValidateTensor("input", input); + if (status != INFINI_OPS_STATUS_SUCCESS) { + return status; + } + status = ValidateTensor("other", other); + if (status != INFINI_OPS_STATUS_SUCCESS) { + return status; + } + status = ValidateTensor("output", output); + if (status != INFINI_OPS_STATUS_SUCCESS) { + return status; + } + + const infini::ops::Handle default_handle; + const infini::ops::Config default_config; + infini::ops::functional::Add( + handle == nullptr ? default_handle : handle->handle, + config == nullptr ? default_config : config->config, + ToInternalTensor(*input), ToInternalTensor(*other), + ToInternalTensor(*output)); + SetLastError(""); + return INFINI_OPS_STATUS_SUCCESS; + } catch (const std::bad_alloc &) { + SetLastError("out of memory while running `infiniOpsAdd`"); + return INFINI_OPS_STATUS_OUT_OF_MEMORY; + } catch (const std::exception &error) { + SetLastError(error.what()); + return INFINI_OPS_STATUS_INTERNAL_ERROR; + } catch (...) { + SetLastError("unknown error while running `infiniOpsAdd`"); + return INFINI_OPS_STATUS_INTERNAL_ERROR; + } +} + +} // extern "C" diff --git a/src/infini/ops.map b/src/infini/ops.map new file mode 100644 index 00000000..1c3c7511 --- /dev/null +++ b/src/infini/ops.map @@ -0,0 +1,10 @@ +INFINIOPS_1 { + global: + extern "C++" { + infini::ops::*; + infini::ops::functional::*; + }; + infiniOps*; + local: + *; +}; diff --git a/tests/test_c_api.py b/tests/test_c_api.py new file mode 100644 index 00000000..265da45b --- /dev/null +++ b/tests/test_c_api.py @@ -0,0 +1,208 @@ +import os +import subprocess +import textwrap +from pathlib import Path + +import pytest + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def test_c_api_header_compiles_with_c(tmp_path): + source = tmp_path / "header_smoke.c" + source.write_text( + '#include \n' + "int main(void) { return INFINI_OPS_STATUS_SUCCESS; }\n" + ) + output = tmp_path / "header_smoke.o" + + _run( + [ + _compiler("CC", "cc"), + "-std=c11", + "-Werror", + *_include_flags(), + "-c", + str(source), + "-o", + str(output), + ] + ) + + +def test_c_api_header_compiles_with_cpp(tmp_path): + source = tmp_path / "header_smoke.cc" + source.write_text( + '#include \n' + "int main() { return INFINI_OPS_STATUS_SUCCESS; }\n" + ) + output = tmp_path / "header_smoke.o" + + _run( + [ + _compiler("CXX", "c++"), + "-std=c++17", + "-Werror", + *_include_flags(require_cpp_api=True), + "-c", + str(source), + "-o", + str(output), + ] + ) + + +def test_c_api_add_smoke(tmp_path): + library_dir = _installed_library_dir() + source = tmp_path / "add_smoke.c" + binary = tmp_path / "add_smoke" + source.write_text(_ADD_SMOKE_SOURCE) + + _run( + [ + _compiler("CC", "cc"), + "-std=c11", + "-Werror", + *_include_flags(), + str(source), + f"-L{library_dir}", + "-linfiniops", + f"-Wl,-rpath,{library_dir}", + "-o", + str(binary), + ] + ) + _run([str(binary)]) + + +def _compiler(env_name, default): + compiler = os.environ.get(env_name, default) + + if not compiler: + pytest.skip(f"`{env_name}` is not configured.") + + return compiler + + +def _include_flags(require_cpp_api=False): + install_prefix = os.environ.get("INFINIOPS_INSTALL_PREFIX") + + if install_prefix: + return [f"-I{Path(install_prefix) / 'include'}"] + + flags = [f"-I{PROJECT_ROOT / 'include'}"] + generated_include_dir = PROJECT_ROOT / "generated" / "include" + + if generated_include_dir.exists(): + flags.append(f"-I{generated_include_dir}") + elif require_cpp_api: + pytest.skip("generated C++ API headers are not available.") + + return flags + + +def _installed_library_dir(): + install_prefix = os.environ.get("INFINIOPS_INSTALL_PREFIX") + + if install_prefix: + for name in ("lib", "lib64"): + library_dir = Path(install_prefix) / name + if (library_dir / "libinfiniops.so").exists(): + return library_dir + pytest.skip(f"`libinfiniops.so` was not found under `{install_prefix}`.") + + library_dir = os.environ.get("INFINIOPS_LIBRARY_DIR") + + if library_dir: + return Path(library_dir) + + try: + import infini.ops + except ImportError as error: + pytest.skip( + "`infini.ops` is not installed and neither " + "`INFINIOPS_INSTALL_PREFIX` nor `INFINIOPS_LIBRARY_DIR` is set: " + f"{error}" + ) + + return Path(infini.ops.__file__).resolve().parent + + +def _run(command): + try: + subprocess.run(command, check=True, text=True, capture_output=True) + except FileNotFoundError as error: + pytest.skip(f"`{command[0]}` is not available: {error}") + except subprocess.CalledProcessError as error: + output = "\n".join((error.stdout, error.stderr)).strip() + raise AssertionError(output) from error + + +_ADD_SMOKE_SOURCE = textwrap.dedent( + r""" + #include + + #include + #include + + int main(void) { + int64_t shape[1] = {3}; + + float input_data[3] = {1.0f, 2.0f, 3.0f}; + float other_data[3] = {4.0f, 5.0f, 6.0f}; + float output_data[3] = {0.0f, 0.0f, 0.0f}; + + InfiniOpsTensor input = {0}; + input.structure_size = sizeof(input); + input.data = input_data; + input.byte_size = sizeof(input_data); + input.data_type = INFINI_OPS_DATA_TYPE_FLOAT32; + input.device_type = INFINI_OPS_DEVICE_TYPE_CPU; + input.rank = 1; + input.shape = shape; + + InfiniOpsTensor other = input; + other.data = other_data; + other.byte_size = sizeof(other_data); + + InfiniOpsTensor output = input; + output.data = output_data; + output.byte_size = sizeof(output_data); + + InfiniOpsHandleAttributes handle_attributes = {0}; + handle_attributes.structure_size = sizeof(handle_attributes); + InfiniOpsHandle handle = NULL; + if (infiniOpsCreateHandle(&handle_attributes, &handle) != + INFINI_OPS_STATUS_SUCCESS) { + return 1; + } + + InfiniOpsConfigAttributes config_attributes = {0}; + config_attributes.structure_size = sizeof(config_attributes); + InfiniOpsConfig config = NULL; + if (infiniOpsCreateConfig(&config_attributes, &config) != + INFINI_OPS_STATUS_SUCCESS) { + return 2; + } + + if (infiniOpsAdd(handle, config, &input, &other, &output) != + INFINI_OPS_STATUS_SUCCESS) { + return 3; + } + + if (output_data[0] != 5.0f || output_data[1] != 7.0f || + output_data[2] != 9.0f) { + return 4; + } + + if (infiniOpsDestroyConfig(config) != INFINI_OPS_STATUS_SUCCESS) { + return 5; + } + if (infiniOpsDestroyHandle(handle) != INFINI_OPS_STATUS_SUCCESS) { + return 6; + } + return 0; + } + """ +).lstrip()