# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0

import _cython_3_2_4
import enum
from typing import Any, Callable, ClassVar

__pyx_capi__: dict
__test__: dict
buffer_deregister: _cython_3_2_4.cython_function_or_method
buffer_register: _cython_3_2_4.cython_function_or_method
check_status: _cython_3_2_4.cython_function_or_method
create: _cython_3_2_4.cython_function_or_method
destroy: _cython_3_2_4.cython_function_or_method
free: _cython_3_2_4.cython_function_or_method
geadd: _cython_3_2_4.cython_function_or_method
geadd_buffer_size: _cython_3_2_4.cython_function_or_method
gemm: _cython_3_2_4.cython_function_or_method
gemm_buffer_size: _cython_3_2_4.cython_function_or_method
gemr2d: _cython_3_2_4.cython_function_or_method
gemr2d_buffer_size: _cython_3_2_4.cython_function_or_method
get_emulation_strategy: _cython_3_2_4.cython_function_or_method
get_matmul_descriptor_attribute_dtype: _cython_3_2_4.cython_function_or_method
get_status_string: _cython_3_2_4.cython_function_or_method
get_stream: _cython_3_2_4.cython_function_or_method
get_version: _cython_3_2_4.cython_function_or_method
grid_create: _cython_3_2_4.cython_function_or_method
grid_destroy: _cython_3_2_4.cython_function_or_method
malloc: _cython_3_2_4.cython_function_or_method
matmul: _cython_3_2_4.cython_function_or_method
matmul_buffer_size: _cython_3_2_4.cython_function_or_method
matmul_descriptor_create: _cython_3_2_4.cython_function_or_method
matmul_descriptor_destroy: _cython_3_2_4.cython_function_or_method
matmul_descriptor_get_attribute: _cython_3_2_4.cython_function_or_method
matmul_descriptor_init: _cython_3_2_4.cython_function_or_method
matmul_descriptor_set_attribute: _cython_3_2_4.cython_function_or_method
matrix_descriptor_create: _cython_3_2_4.cython_function_or_method
matrix_descriptor_destroy: _cython_3_2_4.cython_function_or_method
matrix_descriptor_init: _cython_3_2_4.cython_function_or_method
numroc: _cython_3_2_4.cython_function_or_method
set_emulation_strategy: _cython_3_2_4.cython_function_or_method
set_stream: _cython_3_2_4.cython_function_or_method
syrk: _cython_3_2_4.cython_function_or_method
syrk_buffer_size: _cython_3_2_4.cython_function_or_method
tradd: _cython_3_2_4.cython_function_or_method
tradd_buffer_size: _cython_3_2_4.cython_function_or_method
trmr2d: _cython_3_2_4.cython_function_or_method
trmr2d_buffer_size: _cython_3_2_4.cython_function_or_method
trsm: _cython_3_2_4.cython_function_or_method
trsm_buffer_size: _cython_3_2_4.cython_function_or_method

class EmulationStrategy(enum.IntEnum):
    """
    See `cublasMpEmulationStrategy_t`.
    """
    __new__: ClassVar[Callable] = ...
    DEFAULT: ClassVar[EmulationStrategy] = ...
    EAGER: ClassVar[EmulationStrategy] = ...
    PERFORMANT: ClassVar[EmulationStrategy] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class GridLayout(enum.IntEnum):
    """
    See `cublasMpGridLayout_t`.
    """
    __new__: ClassVar[Callable] = ...
    COL_MAJOR: ClassVar[GridLayout] = ...
    ROW_MAJOR: ClassVar[GridLayout] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class MatmulAlgoType(enum.IntEnum):
    """
    See `cublasMpMatmulAlgoType_t`.
    """
    __new__: ClassVar[Callable] = ...
    ATOMIC_MULTICAST: ClassVar[MatmulAlgoType] = ...
    ATOMIC_P2P: ClassVar[MatmulAlgoType] = ...
    DEFAULT: ClassVar[MatmulAlgoType] = ...
    NO_OVERLAP: ClassVar[MatmulAlgoType] = ...
    SPLIT_MULTICAST: ClassVar[MatmulAlgoType] = ...
    SPLIT_P2P: ClassVar[MatmulAlgoType] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class MatmulDescriptorAttribute(enum.IntEnum):
    """
    See `cublasMpMatmulDescriptorAttribute_t`.
    """
    __new__: ClassVar[Callable] = ...
    ALGO_TYPE: ClassVar[MatmulDescriptorAttribute] = ...
    AMAX_D_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    A_SCALE_MODE: ClassVar[MatmulDescriptorAttribute] = ...
    A_SCALE_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    BIAS_BATCH_STRIDE: ClassVar[MatmulDescriptorAttribute] = ...
    BIAS_DATA_TYPE: ClassVar[MatmulDescriptorAttribute] = ...
    BIAS_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    BIAS_RESULT_SCHEME: ClassVar[MatmulDescriptorAttribute] = ...
    B_SCALE_MODE: ClassVar[MatmulDescriptorAttribute] = ...
    B_SCALE_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    COMMUNICATION_SM_COUNT: ClassVar[MatmulDescriptorAttribute] = ...
    COMMUNICATION_TYPE: ClassVar[MatmulDescriptorAttribute] = ...
    COMPUTE_TYPE: ClassVar[MatmulDescriptorAttribute] = ...
    C_SCALE_MODE: ClassVar[MatmulDescriptorAttribute] = ...
    C_SCALE_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    D_OUT_SCALE_MODE: ClassVar[MatmulDescriptorAttribute] = ...
    D_OUT_SCALE_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    D_SCALE_MODE: ClassVar[MatmulDescriptorAttribute] = ...
    D_SCALE_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_AMAX_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_BATCH_STRIDE: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_DATA_TYPE: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_LD: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_SCALE_MODE: ClassVar[MatmulDescriptorAttribute] = ...
    EPILOGUE_AUX_SCALE_POINTER: ClassVar[MatmulDescriptorAttribute] = ...
    TRANSA: ClassVar[MatmulDescriptorAttribute] = ...
    TRANSB: ClassVar[MatmulDescriptorAttribute] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class MatmulEpilogue(enum.IntEnum):
    """
    See `cublasMpMatmulEpilogue_t`.
    """
    __new__: ClassVar[Callable] = ...
    ALLREDUCE: ClassVar[MatmulEpilogue] = ...
    BGRADA: ClassVar[MatmulEpilogue] = ...
    BGRADB: ClassVar[MatmulEpilogue] = ...
    BIAS: ClassVar[MatmulEpilogue] = ...
    DEFAULT: ClassVar[MatmulEpilogue] = ...
    DGELU: ClassVar[MatmulEpilogue] = ...
    DGELU_BGRAD: ClassVar[MatmulEpilogue] = ...
    DRELU: ClassVar[MatmulEpilogue] = ...
    DRELU_BGRAD: ClassVar[MatmulEpilogue] = ...
    GELU: ClassVar[MatmulEpilogue] = ...
    GELU_AUX: ClassVar[MatmulEpilogue] = ...
    GELU_AUX_BIAS: ClassVar[MatmulEpilogue] = ...
    GELU_BIAS: ClassVar[MatmulEpilogue] = ...
    RELU: ClassVar[MatmulEpilogue] = ...
    RELU_AUX: ClassVar[MatmulEpilogue] = ...
    RELU_AUX_BIAS: ClassVar[MatmulEpilogue] = ...
    RELU_BIAS: ClassVar[MatmulEpilogue] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class MatmulMatrixScale(enum.IntEnum):
    """
    See `cublasMpMatmulMatrixScale_t`.
    """
    __new__: ClassVar[Callable] = ...
    BLK128x128_FP32: ClassVar[MatmulMatrixScale] = ...
    OUTER_VEC_FP32: ClassVar[MatmulMatrixScale] = ...
    SCALAR_FP32: ClassVar[MatmulMatrixScale] = ...
    VEC128_FP32: ClassVar[MatmulMatrixScale] = ...
    VEC16_UE4M3: ClassVar[MatmulMatrixScale] = ...
    VEC32_UE8M0: ClassVar[MatmulMatrixScale] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class ResultScheme(enum.IntEnum):
    """
    See `cublasMpResultScheme_t`.
    """
    __new__: ClassVar[Callable] = ...
    DEFAULT: ClassVar[ResultScheme] = ...
    FULL: ClassVar[ResultScheme] = ...
    PARTIAL: ClassVar[ResultScheme] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class Status(enum.IntEnum):
    """
    See `cublasMpStatus_t`.
    """
    __new__: ClassVar[Callable] = ...
    ALLOCATION_FAILED: ClassVar[Status] = ...
    ARCHITECTURE_MISMATCH: ClassVar[Status] = ...
    EXECUTION_FAILED: ClassVar[Status] = ...
    INTERNAL_ERROR: ClassVar[Status] = ...
    INVALID_VALUE: ClassVar[Status] = ...
    NOT_INITIALIZED: ClassVar[Status] = ...
    NOT_SUPPORTED: ClassVar[Status] = ...
    SUCCESS: ClassVar[Status] = ...
    _generate_next_value_: ClassVar[Callable] = ...
    _member_map_: ClassVar[dict] = ...
    _member_names_: ClassVar[list] = ...
    _member_type_: ClassVar[type[int]] = ...
    _value2member_map_: ClassVar[dict] = ...

class cuBLASMpError(Exception):
    def __init__(self, status) -> Any:
        """cuBLASMpError.__init__(self, status)"""
    def __reduce__(self) -> Any:
        """cuBLASMpError.__reduce__(self)"""
