diff --git a/.gitignore b/.gitignore index 84244ce82..bbe6c5292 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ *.gcda *.gcov core +!core/ *.fluid *.pyc *.swp diff --git a/conda/conda-build/meta.yaml b/conda/conda-build/meta.yaml index 85ea88d74..92e9acfab 100644 --- a/conda/conda-build/meta.yaml +++ b/conda/conda-build/meta.yaml @@ -144,6 +144,7 @@ requirements: - libnvjitlink - libcusparse {% endif %} + - numba >=0.57.1 - opt_einsum >=3.3 - scipy - typing_extensions diff --git a/continuous_integration/scripts/test-cunumeric b/continuous_integration/scripts/test-cunumeric index 698179b31..605583f34 100755 --- a/continuous_integration/scripts/test-cunumeric +++ b/continuous_integration/scripts/test-cunumeric @@ -1,7 +1,7 @@ #!/usr/bin/env bash setup_env() { - mamba create -yn legate -c ~/.artifacts/conda-build/legate_core -c ~/.artifacts/conda-build/cunumeric -c conda-forge -c "nvidia/label/cuda-12.0.0" legate-core cunumeric + mamba create -yn legate -c ~/.artifacts/conda-build/legate_core -c ~/.artifacts/conda-build/cunumeric -c conda-forge -c "nvidia/label/cuda-12.0.0" legate-core cunumeric cuda-nvcc } setup_test_env() { diff --git a/cunumeric/_sphinxext/_cunumeric_directive.py b/cunumeric/_sphinxext/_cunumeric_directive.py index ef6402f6c..03fb3be32 100644 --- a/cunumeric/_sphinxext/_cunumeric_directive.py +++ b/cunumeric/_sphinxext/_cunumeric_directive.py @@ -15,14 +15,14 @@ from __future__ import annotations from docutils import nodes -from docutils.statemachine import ViewList +from docutils.statemachine import StringList from sphinx.util.docutils import SphinxDirective from sphinx.util.nodes import nested_parse_with_titles class CunumericDirective(SphinxDirective): def parse(self, rst_text: str, annotation: str) -> list[nodes.Node]: - result = ViewList() + result = StringList() for line in rst_text.split("\n"): result.append(line, annotation) node = nodes.paragraph() diff --git a/cunumeric/module.py b/cunumeric/module.py index 424f89df4..e87670b75 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -1748,7 +1748,7 @@ def check_list_depth(arr: Any, prefix: NdShape = (0,)) -> int: "List depths are mismatched. First element was at depth " f"{first_depth}, but there is an element at" f" depth {other_depth}, " - f"arrays{convert_to_array_form(prefix+(idx+1,))}" + f"arrays{convert_to_array_form(prefix + (idx + 1,))}" ) return depths[0] + 1 diff --git a/cunumeric/numba_utils.py b/cunumeric/numba_utils.py new file mode 100644 index 000000000..4aa46cfab --- /dev/null +++ b/cunumeric/numba_utils.py @@ -0,0 +1,256 @@ +# Copyright 2023 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + Optional, + Tuple, + Union, + cast, +) + +from llvmlite import ir +from llvmlite.ir.builder import IRBuilder +from llvmlite.ir.instructions import Ret +from llvmlite.ir.types import FunctionType +from llvmlite.ir.values import Value +from numba import types +from numba.core import sigutils +from numba.core.base import BaseContext +from numba.core.callconv import BaseCallConv +from numba.core.codegen import CodeLibrary +from numba.core.compiler_lock import global_compiler_lock +from numba.core.funcdesc import FunctionDescriptor +from numba.core.typing.templates import Signature +from numba.cuda.codegen import CUDACodeLibrary +from numba.cuda.compiler import compile_cuda + + +class SoACallConv(BaseCallConv): + """ + Calling convention where returned values are stored through pointers + provided as arguments. + + - If the return type is a scalar, the first argument is a pointer to the + return type. + - If the return type is a tuple of length N, then the first N arguments are + pointers to each of the elements of the tuple. + + In equivalent C, the prototype of a function with this calling convention + would take the following form: + + void (*, ..., *, + ); + """ + + def _make_call_helper(self, builder: Any) -> None: + # Call helpers are used for the exception implementation. This is not + # needed when only wrapping functions. + msg = "Python exceptions are unsupported when returning in SoA form" + raise NotImplementedError(msg) + + def return_value(self, builder: IRBuilder, retval: Value) -> Ret: + return builder.ret(retval) + + def return_user_exc( + self, + builder: IRBuilder, + exc: Any, + exc_args: Any = None, + loc: Any = None, + func_name: Any = None, + ) -> None: + msg = "Python exceptions are unsupported when returning in SoA form" + raise NotImplementedError(msg) + + def return_status_propagate(self, builder: IRBuilder, status: Any) -> None: + msg = "Return status is unsupported when returning in SoA form" + raise NotImplementedError(msg) + + def get_function_type( + self, restype: types.Type, argtypes: Iterable[types.Type] + ) -> FunctionType: + """ + Get the LLVM IR Function type for *restype* and *argtypes*. + """ + arginfo = self._get_arg_packer(argtypes) + be_argtypes = list(arginfo.argument_types) + if isinstance(restype, types.BaseTuple): + return_types = [self.get_return_type(t) for t in restype.types] + else: + return_types = [self.get_return_type(restype)] + fnty = ir.FunctionType(ir.VoidType(), return_types + be_argtypes) + return fnty + + def decorate_function( + self, + fn: Callable[[Any], Any], + args: Iterable[str], + fe_argtypes: List[types.Type], + noalias: bool = False, + ) -> None: + """ + Set names and attributes of function arguments. + """ + raise NotImplementedError("Function decoration not used for SoA ABI") + + def get_arguments( + self, func: ir.Function, restype: types.Type + ) -> Tuple[ir.Argument, ...]: + """ + Get the Python-level arguments of LLVM *func*. + """ + if isinstance(restype, types.BaseTuple): + n_returns = len(restype.types) + else: + n_returns = 1 + + return func.args[n_returns:] + + def call_function( + self, + builder: ir.IRBuilder, + callee: ir.Function, + resty: types.Type, + argtys: Iterable[types.Type], + args: Iterable[ir.Value], + attrs: Optional[Tuple[str, ...]] = None, + ) -> Tuple[ir.Value, ir.Value]: + """ + Call the Numba-compiled *callee*. + """ + raise NotImplementedError("Can't call SoA return function directly") + + +def soa_wrap_function( + context: BaseContext, + lib: CodeLibrary, + fndesc: FunctionDescriptor, + nvvm_options: Dict[str, Union[int, str, None]], + wrapper_name: str, +) -> CUDACodeLibrary: + """ + Wrap a Numba ABI function such that it returns tuple values into SoA + arguments. + """ + new_library = lib.codegen.create_library( + f"{lib.name}_function_", + entry_name=wrapper_name, + nvvm_options=nvvm_options, + ) + library = cast(CUDACodeLibrary, new_library) + library.add_linking_library(lib) + + # Determine the caller (C ABI) and wrapper (Numba ABI) function types + argtypes = fndesc.argtypes + restype = fndesc.restype + soa_call_conv = SoACallConv(context) + wrapperty = soa_call_conv.get_function_type(restype, argtypes) + calleety = context.call_conv.get_function_type(restype, argtypes) + + # Create a new module and declare the callee + wrapper_module = context.create_module("cuda.soa.wrapper") + callee = ir.Function(wrapper_module, calleety, fndesc.llvm_func_name) + + # Define the caller - populate it with a call to the callee and return + # its return value + + wrapper = ir.Function(wrapper_module, wrapperty, wrapper_name) + builder = ir.IRBuilder(wrapper.append_basic_block("")) + + arginfo = context.get_arg_packer(argtypes) + wrapper_args = soa_call_conv.get_arguments(wrapper, restype) + callargs = arginfo.as_arguments(builder, wrapper_args) + # We get (status, return_value), but we ignore the status since we + # can't propagate it through the SoA ABI anyway + _, return_value = context.call_conv.call_function( + builder, callee, restype, argtypes, callargs + ) + + if isinstance(restype, types.BaseTuple): + for i in range(len(restype.types)): + val = builder.extract_value(return_value, i) + builder.store(val, wrapper.args[i]) + else: + builder.store(return_value, wrapper.args[0]) + builder.ret_void() + + library.add_ir_module(wrapper_module) + library.finalize() + return library + + +@global_compiler_lock +def compile_ptx_soa( + pyfunc: Callable[..., Any], + sig: Union[Tuple[types.Type], str, Signature], + debug: bool = False, + lineinfo: bool = False, + device: bool = False, + fastmath: bool = False, + cc: Optional[Tuple[int, int]] = None, + opt: bool = True, + abi_info: Optional[Dict[str, str]] = None, +) -> Tuple[str, types.Type]: + # This is just a copy of Numba's compile_ptx, with a modification to return + # values as SoA and some simplifications to keep it short + if not device: + raise NotImplementedError( + "Only device functions can be compiled for the SoA ABI" + ) + + nvvm_options: Dict[str, Union[int, str, None]] = { + "fastmath": fastmath, + "opt": 3 if opt else 0, + } + + # Use the Python function name as the function name in PTX if it is not + # specified - otherwise, use the specified name. + if abi_info: + wrapper_name = abi_info["abi_name"] + else: + wrapper_name = pyfunc.__name__ + + args, return_type = sigutils.normalize_signature(sig) + + # Default to Compute Capability 5.0 if not specified + cc = cc or (5, 0) + + cres = compile_cuda( + pyfunc, + return_type, + args, + debug=debug, + lineinfo=lineinfo, + fastmath=fastmath, + nvvm_options=nvvm_options, + cc=cc, + ) + + lib = soa_wrap_function( + cres.target_context, + cres.library, + cres.fndesc, + nvvm_options, + wrapper_name, + ) + + ptx = lib.get_asm_str(cc=cc) + resty = cres.signature.return_type + + return ptx, resty diff --git a/install.py b/install.py index 92d16ad55..b35eb3681 100755 --- a/install.py +++ b/install.py @@ -336,10 +336,15 @@ def validate_path(path): if debug or verbose: cmake_flags += ["--log-level=%s" % ("DEBUG" if debug else "VERBOSE")] + if debug: + build_type = "Debug" + elif debug_release: + build_type = "RelWithDebInfo" + else: + build_type = "Release" + cmake_flags += f"""\ --DCMAKE_BUILD_TYPE={( - "Debug" if debug else "RelWithDebInfo" if debug_release else "Release" -)} +-DCMAKE_BUILD_TYPE={build_type} -DBUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES={str(arch)} -DLegion_MAX_DIM={str(maxdim)} diff --git a/pyproject.toml b/pyproject.toml index 04c31fb7a..c766329ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,12 +39,14 @@ exclude = ''' _build | buck-out | build | - dist + dist | + typings )/ ''' [tool.mypy] python_version = "3.10" +mypy_path = "typings/" pretty = true show_error_codes = true diff --git a/tests/integration/utils/comparisons.py b/tests/integration/utils/comparisons.py index 65571b38c..9ab5247e9 100644 --- a/tests/integration/utils/comparisons.py +++ b/tests/integration/utils/comparisons.py @@ -50,7 +50,7 @@ def allclose( inds = islice(zip(*np.where(~close)), diff_limit) diffs = [f" index {i}: {a[i]} {b[i]}" for i in inds] N = len(diffs) - print(f"First {N} difference{'s' if N>1 else ''} for allclose:\n") + print(f"First {N} difference{'s' if N > 1 else ''} for allclose:\n") print("\n".join(diffs)) print(f"\nWith diff_limit={diff_limit}\n") diff --git a/tests/unit/cunumeric/test_numba_utils.py b/tests/unit/cunumeric/test_numba_utils.py new file mode 100644 index 000000000..269be629c --- /dev/null +++ b/tests/unit/cunumeric/test_numba_utils.py @@ -0,0 +1,208 @@ +# Copyright 2023 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import re +import typing + +import pytest +from numba.types import Tuple, UniTuple, float32, float64, int32, int64 + +from cunumeric.numba_utils import compile_ptx_soa + + +@pytest.fixture +def addsub() -> typing.Callable: + def addsub(x, y): + return x + y, x - y + + return addsub + + +@pytest.fixture +def addsubmul() -> typing.Callable: + def addsubmul(a, x, y): + return a * (x + y), a * (x - y) + + return addsubmul + + +@pytest.fixture +def addsubconst() -> typing.Callable: + def addsubconst(x, y): + return x + y, x - y, 3 + + return addsubconst + + +def param_pattern( + function_name: str, param_index: int, param_type: str, reg_prefix: str +) -> str: + """ + A helper function to generate patterns for recognizing parameter references + in PTX functions - an example of a parameter reference looks like: + + ld.param.u64 %rd1, [addsubconst_param_0]; + + These usually appear at the beginning of a function. + """ + return ( + rf"ld\.param\.{param_type}" + rf"\s+%{reg_prefix}[0-9]+,\s+" + rf"\[{function_name}_param_" + rf"{param_index}\]" + ) + + +def test_soa(addsub) -> None: + # A basic test of compilation with an SoA interface + + # Compile for two int32 inputs and two int32 outputs + signature = UniTuple(int32, 2)(int32, int32) + ptx, resty = compile_ptx_soa(addsub, signature, device=True) + + # The function definition should use the name of the Python function + fn_def_pattern = r"\.visible\s+\.func\s+addsub" + assert re.search(fn_def_pattern, ptx) + + # The return type should match that of the signature's return type + assert resty == signature.return_type + + # The function should have 4 parameters (numbered 0 to 3) + assert re.search("addsub_param_3", ptx) + assert not re.search("addsub_param_4", ptx) + + # The first two parameters should be treated as pointers (u64 values) + assert re.search(param_pattern("addsub", 0, "u64", "rd"), ptx) + assert re.search(param_pattern("addsub", 1, "u64", "rd"), ptx) + + # The remaining two parameters should be treated as 32 bit integers + assert re.search(param_pattern("addsub", 2, "u32", "r"), ptx) + assert re.search(param_pattern("addsub", 3, "u32", "r"), ptx) + + +def test_soa_fn_name(addsub) -> None: + # Ensure that when a wrapper function name is specified, it is used in the + # PTX. + signature = UniTuple(int32, 2)(int32, int32) + abi_info = {"abi_name": "addsub_soa"} + ptx, resty = compile_ptx_soa( + addsub, signature, device=True, abi_info=abi_info + ) + fn_def_pattern = r"\.visible\s+\.func\s+addsub_soa" + assert re.search(fn_def_pattern, ptx) + + +def test_soa_arg_types(addsub) -> None: + # Ensure that specifying a different argument type is reflected + # appropriately in the generated PTX + signature = UniTuple(int32, 2)(int32, int64) + ptx, resty = compile_ptx_soa(addsub, signature, device=True) + + # The final two parameters should now be a 32- and a 64-bit values + # respectively. Note that the load of the last parameter may be an + # instruction with a 32-bit destination type that effectively chops off the + # upper 32 bits, so we cannot test for a load of a 64-bit value, which + # would look like: + # + # ld.param.u64 %rd2, [addsub_param_3]; + # + # but instead we'd potentially get + # + # ld.param.u32 %r2, [addsub_param_3]; + # + # So we test the bit width of the parameters only: + assert re.search(r".param\s+.b32\s+addsub_param_2", ptx) + assert re.search(r".param\s+.b64\s+addsub_param_3", ptx) + + +def test_soa_more_args(addsubmul) -> None: + # A test with three arguments, but only two return values + + signature = UniTuple(int32, 2)(int32, int32, int32) + ptx, resty = compile_ptx_soa(addsubmul, signature, device=True) + + # The function should have 5 parameters (numbered 0 to 4) + assert re.search("addsubmul_param_4", ptx) + assert not re.search("addsubmul_param_5", ptx) + + # The first two parameters should be treated as pointers (u64 values) + assert re.search(param_pattern("addsubmul", 0, "u64", "rd"), ptx) + assert re.search(param_pattern("addsubmul", 1, "u64", "rd"), ptx) + + # The remaining three parameters should be treated as 32 bit integers + assert re.search(param_pattern("addsubmul", 2, "u32", "r"), ptx) + assert re.search(param_pattern("addsubmul", 3, "u32", "r"), ptx) + assert re.search(param_pattern("addsubmul", 4, "u32", "r"), ptx) + + +def test_soa_more_returns(addsubconst) -> None: + # Test with two arguments and three return values + + signature = UniTuple(int32, 3)(int32, int32) + ptx, resty = compile_ptx_soa(addsubconst, signature, device=True) + + # The function should have 5 parameters (numbered 0 to 4) + assert re.search("addsubconst_param_4", ptx) + assert not re.search("addsubconst_param_5", ptx) + + # The first three parameters should be treated as pointers (u64 values) + assert re.search(param_pattern("addsubconst", 0, "u64", "rd"), ptx) + assert re.search(param_pattern("addsubconst", 1, "u64", "rd"), ptx) + assert re.search(param_pattern("addsubconst", 2, "u64", "rd"), ptx) + + # The remaining two parameters should be treated as 32 bit integers + assert re.search(param_pattern("addsubconst", 3, "u32", "r"), ptx) + assert re.search(param_pattern("addsubconst", 4, "u32", "r"), ptx) + + +def test_soa_varying_types(addsub) -> None: + # Argument types differ from each other and the return type + + signature = UniTuple(float64, 2)(int32, float32) + ptx, resty = compile_ptx_soa(addsub, signature, device=True) + + # The first two parameters should be treated as pointers (u64 values) + assert re.search(param_pattern("addsub", 0, "u64", "rd"), ptx) + assert re.search(param_pattern("addsub", 1, "u64", "rd"), ptx) + + # The remaining two parameters should be a 32-bit integer and a 32-bit + # float + assert re.search(param_pattern("addsub", 2, "u32", "r"), ptx) + assert re.search(param_pattern("addsub", 3, "f32", "f"), ptx) + + # There should be a 64-bit floating point store for the result + assert re.search(r"st\.f64", ptx) + + +def test_soa_heterogeneous_return_type(addsubconst) -> None: + # Test with return values of different types + + signature = Tuple((float32, float64, int32))(float32, float32) + ptx, resty = compile_ptx_soa(addsubconst, signature, device=True) + + # There should be 32- and 64-bit floating point, and 32-bit integer stores + # for the result + assert re.search(r"st\.f32", ptx) + assert re.search(r"st\.f64", ptx) + assert re.search(r"st\.u32", ptx) + + +# Test of one return value + +# Test of not putting device in + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(sys.argv)) diff --git a/typings/llvmlite/__init__.pyi b/typings/llvmlite/__init__.pyi new file mode 100644 index 000000000..711280d9e --- /dev/null +++ b/typings/llvmlite/__init__.pyi @@ -0,0 +1,15 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from llvmlite import ir diff --git a/typings/llvmlite/ir/__init__.pyi b/typings/llvmlite/ir/__init__.pyi new file mode 100644 index 000000000..048ddc711 --- /dev/null +++ b/typings/llvmlite/ir/__init__.pyi @@ -0,0 +1,18 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from llvmlite.ir.builder import * +from llvmlite.ir.module import * +from llvmlite.ir.types import * +from llvmlite.ir.values import * diff --git a/typings/llvmlite/ir/_utils.pyi b/typings/llvmlite/ir/_utils.pyi new file mode 100644 index 000000000..d02b145ad --- /dev/null +++ b/typings/llvmlite/ir/_utils.pyi @@ -0,0 +1,17 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +class _HasMetadata: ... +class _StrCaching: ... +class _StringReferenceCaching: ... diff --git a/typings/llvmlite/ir/builder.pyi b/typings/llvmlite/ir/builder.pyi new file mode 100644 index 000000000..29b4a8959 --- /dev/null +++ b/typings/llvmlite/ir/builder.pyi @@ -0,0 +1,32 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Iterable, Optional, Union + +from llvmlite.ir.instructions import Instruction, Ret +from llvmlite.ir.values import Block, Value + +class IRBuilder: + def __init__(self, block: Optional[Block]): ... + def ret(self, return_value: Value) -> Ret: ... + def extract_value( + self, + agg: Value, + idx: Union[Iterable[int], int], + name: str = "", + ) -> Instruction: ... + def store( + self, value: Value, ptr: Value, align: Optional[int] = None + ) -> Instruction: ... + def ret_void(self) -> Ret: ... diff --git a/typings/llvmlite/ir/instructions.pyi b/typings/llvmlite/ir/instructions.pyi new file mode 100644 index 000000000..75f1c2872 --- /dev/null +++ b/typings/llvmlite/ir/instructions.pyi @@ -0,0 +1,20 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from llvmlite.ir._utils import _HasMetadata +from llvmlite.ir.values import NamedValue + +class Instruction(NamedValue, _HasMetadata): ... +class Terminator(Instruction): ... +class Ret(Terminator): ... diff --git a/typings/llvmlite/ir/module.pyi b/typings/llvmlite/ir/module.pyi new file mode 100644 index 000000000..e9d24b059 --- /dev/null +++ b/typings/llvmlite/ir/module.pyi @@ -0,0 +1,15 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +class Module: ... diff --git a/typings/llvmlite/ir/types.pyi b/typings/llvmlite/ir/types.pyi new file mode 100644 index 000000000..40d78c0b4 --- /dev/null +++ b/typings/llvmlite/ir/types.pyi @@ -0,0 +1,27 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Sequence + +from llvmlite.ir._utils import _StrCaching + +class Type(_StrCaching): ... + +class FunctionType(Type): + def __init__( + self, return_type: Type, args: Sequence[Type], var_arg: bool = False + ): ... + +class PointerType(Type): ... +class VoidType(Type): ... diff --git a/typings/llvmlite/ir/values.pyi b/typings/llvmlite/ir/values.pyi new file mode 100644 index 000000000..6ea7db477 --- /dev/null +++ b/typings/llvmlite/ir/values.pyi @@ -0,0 +1,37 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Tuple + +from llvmlite.ir._utils import ( + _HasMetadata, + _StrCaching, + _StringReferenceCaching, +) +from llvmlite.ir.module import Module +from llvmlite.ir.types import FunctionType + +class Value: ... +class NamedValue(_StrCaching, _StringReferenceCaching, Value): ... +class Block(NamedValue): ... +class _BaseArgument(NamedValue): ... +class Argument(_BaseArgument): ... +class _ConstOpMixin: ... +class GlobalValue(NamedValue, _ConstOpMixin, _HasMetadata): ... + +class Function(GlobalValue): + args: Tuple[Argument] + + def __init__(self, module: Module, ftype: FunctionType, name: str): ... + def append_basic_block(self, name: str) -> Block: ... diff --git a/typings/numba/__init__.pyi b/typings/numba/__init__.pyi new file mode 100644 index 000000000..33afd9e7a --- /dev/null +++ b/typings/numba/__init__.pyi @@ -0,0 +1,25 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Callable + +import numba.cuda # import compile_ptx +from numba.core import types +from numba.core.ccallback import CFunc +from numba.core.types import CPointer, uint64 + +def cfunc(sig: Any) -> Any: + def wrapper(func: Callable[[Any], Any]) -> tuple[Any]: ... + +__all__ = ["types"] diff --git a/typings/numba/core/__init__.pyi b/typings/numba/core/__init__.pyi new file mode 100644 index 000000000..7cc93af29 --- /dev/null +++ b/typings/numba/core/__init__.pyi @@ -0,0 +1,13 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/typings/numba/core/base.pyi b/typings/numba/core/base.pyi new file mode 100644 index 000000000..fdc9e0874 --- /dev/null +++ b/typings/numba/core/base.pyi @@ -0,0 +1,26 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Sequence + +from llvmlite import ir +from numba.core.callconv import BaseCallConv +from numba.core.datamodel import ArgPacker +from numba.core.types import Type + +class BaseContext: + call_conv: BaseCallConv + + def create_module(self, name: str) -> ir.Module: ... + def get_arg_packer(self, fe_args: Sequence[Type]) -> ArgPacker: ... diff --git a/typings/numba/core/callconv.pyi b/typings/numba/core/callconv.pyi new file mode 100644 index 000000000..01d410261 --- /dev/null +++ b/typings/numba/core/callconv.pyi @@ -0,0 +1,39 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Iterable, Optional, Tuple + +from llvmlite.ir.builder import IRBuilder +from llvmlite.ir.types import FunctionType, PointerType, Type as LLType +from llvmlite.ir.values import Function, Value +from numba.core.base import BaseContext +from numba.core.datamodel import ArgPacker +from numba.core.types import Type + +class BaseCallConv: + def __init__(self, context: BaseContext): ... + def _get_arg_packer(self, argtypes: Iterable[Type]) -> ArgPacker: ... + def get_return_type(self, ty: Type) -> LLType: ... + def get_function_type( + self, restype: Type, argtypes: Iterable[Type] + ) -> FunctionType: ... + def call_function( + self, + builder: IRBuilder, + callee: Function, + resty: Type, + argtys: Iterable[Type], + args: Iterable[Value], + attrs: Optional[Tuple[str, ...]] = None, + ) -> Tuple[Value, Value]: ... diff --git a/typings/numba/core/ccallback/__init__.pyi b/typings/numba/core/ccallback/__init__.pyi new file mode 100644 index 000000000..f9b9c57e8 --- /dev/null +++ b/typings/numba/core/ccallback/__init__.pyi @@ -0,0 +1,22 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +class CFunc(object): + def __init__( + self, pyfunc: Any, sig: Any, locals: Any, options: Any + ) -> None: ... + @property + def address(self) -> int: ... diff --git a/typings/numba/core/codegen.pyi b/typings/numba/core/codegen.pyi new file mode 100644 index 000000000..dad46d2ff --- /dev/null +++ b/typings/numba/core/codegen.pyi @@ -0,0 +1,30 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +from llvmlite import ir + +class CodeLibrary: + codegen: "Codegen" + name: str + + def add_linking_library(self, library: CodeLibrary) -> None: ... + def add_ir_module(self, module: ir.Module) -> None: ... + def finalize(self) -> None: ... + +class Codegen: + name: str + + def create_library(self, name: str, **kwargs: Any) -> CodeLibrary: ... diff --git a/typings/numba/core/compiler.pyi b/typings/numba/core/compiler.pyi new file mode 100644 index 000000000..673488778 --- /dev/null +++ b/typings/numba/core/compiler.pyi @@ -0,0 +1,24 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from numba.core.base import BaseContext +from numba.core.codegen import CodeLibrary +from numba.core.funcdesc import FunctionDescriptor +from numba.core.typing.templates import Signature + +class CompileResult: + target_context: BaseContext + library: CodeLibrary + fndesc: FunctionDescriptor + signature: Signature diff --git a/typings/numba/core/compiler_lock.pyi b/typings/numba/core/compiler_lock.pyi new file mode 100644 index 000000000..24616c2e1 --- /dev/null +++ b/typings/numba/core/compiler_lock.pyi @@ -0,0 +1,20 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Callable + +class _CompilerLock: + def __call__(self, func: Callable[..., Any]) -> Callable[[Any], Any]: ... + +global_compiler_lock = _CompilerLock() diff --git a/typings/numba/core/datamodel.pyi b/typings/numba/core/datamodel.pyi new file mode 100644 index 000000000..40f38f97f --- /dev/null +++ b/typings/numba/core/datamodel.pyi @@ -0,0 +1,27 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Sequence, Tuple + +from llvmlite.ir import Argument +from llvmlite.ir.builder import IRBuilder +from llvmlite.ir.types import Type +from llvmlite.ir.values import Value + +class ArgPacker: + argument_types: Sequence[Type] + + def as_arguments( + self, builder: IRBuilder, values: Tuple[Argument, ...] + ) -> Tuple[Value]: ... diff --git a/typings/numba/core/funcdesc.pyi b/typings/numba/core/funcdesc.pyi new file mode 100644 index 000000000..fe83b3d26 --- /dev/null +++ b/typings/numba/core/funcdesc.pyi @@ -0,0 +1,22 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple + +from numba.core.types import Type + +class FunctionDescriptor: + argtypes: Tuple[Type] + restype: Type + llvm_func_name: str diff --git a/typings/numba/core/sigutils.pyi b/typings/numba/core/sigutils.pyi new file mode 100644 index 000000000..2c122b5a7 --- /dev/null +++ b/typings/numba/core/sigutils.pyi @@ -0,0 +1,22 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Tuple, Union + +from numba.core.types import Type +from numba.core.typing.templates import Signature + +def normalize_signature( + sig: Union[Tuple[Type, ...], str, Signature] +) -> Tuple[Tuple[Type, ...], Type]: ... diff --git a/typings/numba/core/types/__init__.pyi b/typings/numba/core/types/__init__.pyi new file mode 100644 index 000000000..9b5cbdfdc --- /dev/null +++ b/typings/numba/core/types/__init__.pyi @@ -0,0 +1,48 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Tuple + +class Opaque: ... + +class NoneType(Opaque): + def __init__(self, name: str) -> None: ... + +class Type: + def __init__(self, name: str) -> None: ... + +class Number(Type): ... + +class Integer(Number): + def __init__(self, name: str) -> None: ... + +class RawPointer: + def __init__(self, name: str) -> None: ... + +class CPointer(Type): + def __init__(self, dtype: Type) -> None: ... + +class Sized(Type): ... +class ConstSized(Type): ... +class Hashable(Type): ... + +class BaseTuple(ConstSized, Hashable): + types: Tuple[Type] + +none = NoneType("none") + +uint32 = Integer("uint32") +uint64 = Integer("uint64") +void = none +voidptr = Type("void*") diff --git a/typings/numba/core/typing/templates.pyi b/typings/numba/core/typing/templates.pyi new file mode 100644 index 000000000..3eac0629a --- /dev/null +++ b/typings/numba/core/typing/templates.pyi @@ -0,0 +1,18 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from numba.core.types import Type + +class Signature: + return_type: Type diff --git a/typings/numba/cuda/__init__.pyi b/typings/numba/cuda/__init__.pyi new file mode 100644 index 000000000..77610ef8c --- /dev/null +++ b/typings/numba/cuda/__init__.pyi @@ -0,0 +1,19 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +from numba.cuda.compiler import compile_ptx as compile_ptx + +def get_current_device() -> Any: ... diff --git a/typings/numba/cuda/codegen.pyi b/typings/numba/cuda/codegen.pyi new file mode 100644 index 000000000..297a07a31 --- /dev/null +++ b/typings/numba/cuda/codegen.pyi @@ -0,0 +1,26 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Optional, Tuple + +from numba.core.codegen import Codegen, CodeLibrary + +class CUDACodeLibrary(CodeLibrary): + codegen: "JITCUDACodegen" + name: str + + def get_asm_str(self, cc: Optional[Tuple[int, int]] = None) -> str: ... + +class JITCUDACodegen(Codegen): + def create_library(self, name: str, **kwargs: Any) -> CUDACodeLibrary: ... diff --git a/typings/numba/cuda/compiler.pyi b/typings/numba/cuda/compiler.pyi new file mode 100644 index 000000000..b7c51bc47 --- /dev/null +++ b/typings/numba/cuda/compiler.pyi @@ -0,0 +1,40 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Callable, Dict, Optional, Tuple, Union + +from numba.core.compiler import CompileResult +from numba.core.types import Type + +def compile_ptx( + pyfunc: Callable[[Any], Any], + args: Any, + debug: bool = False, + lineinfo: bool = False, + device: bool = False, + fastmath: bool = False, + cc: Optional[Any] = None, + opt: bool = True, +) -> tuple[Any]: ... +def compile_cuda( + pyfunc: Callable[[Any], Any], + return_type: Type, + args: Tuple[Type, ...], + debug: bool = False, + lineinfo: bool = False, + inline: bool = False, + fastmath: bool = False, + nvvm_options: Optional[Dict[str, Optional[Union[str, int]]]] = None, + cc: Optional[Tuple[int, int]] = None, +) -> CompileResult: ... diff --git a/typings/numba/cuda/target.pyi b/typings/numba/cuda/target.pyi new file mode 100644 index 000000000..cc43b37b5 --- /dev/null +++ b/typings/numba/cuda/target.pyi @@ -0,0 +1,24 @@ +# Copyright 2024 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from llvmlite import ir +from numba.core.base import BaseContext +from numba.core.callconv import BaseCallConv + +class CUDACallConv(BaseCallConv): ... + +class CUDATargetContext(BaseContext): + call_conv: CUDACallConv + + def create_module(self, name: str) -> ir.Module: ... diff --git a/typings/numba/types/CPointer.pyi b/typings/numba/types/CPointer.pyi new file mode 100644 index 000000000..249a23f19 --- /dev/null +++ b/typings/numba/types/CPointer.pyi @@ -0,0 +1,5 @@ +# import numpy as np +from numba.core.types.abstract import Type + +class CPointer(Type): + def __init__(self, dtype: Type) -> None: ... diff --git a/typings/numba/types/__init__.pyi b/typings/numba/types/__init__.pyi new file mode 100644 index 000000000..14c90eca2 --- /dev/null +++ b/typings/numba/types/__init__.pyi @@ -0,0 +1,12 @@ +class Type: ... +class Number(Type): ... + +class Integer(Number): + def __init__(self, name: str) -> None: ... + +class CPointer(Type): + def __init__(self, dtype: Type) -> None: ... + +uint32 = Integer("uint32") +uint64 = Integer("uint64") +void = None