Skip to content

Vendor in Codegen and CodeLibrary for CUDA-specific customization #327

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gmarkall merged 2 commits into from
Jul 28, 2025
Merged

Vendor in Codegen and CodeLibrary for CUDA-specific customization #327

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d6cb0fc
[Refactor] Vendor in Codegen and CodeLibrary for CUDA-specific changes
VijayKandiah Jul 17, 2025
363a298
[Refactor][NFC] Formatting fixes
VijayKandiah Jul 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion numba_cuda/numba/cuda/codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
from llvmlite import ir

from numba.core import config, serialize
from numba.core.codegen import Codegen, CodeLibrary
from .cudadrv import devices, driver, nvvm, runtime, nvrtc
from numba.cuda.core.codegen import Codegen, CodeLibrary
from numba.cuda.cudadrv.libs import get_cudalib
from numba.cuda.cudadrv.linkable_code import LinkableCode
from numba.cuda.memory_management.nrt import NRT_LIBRARY
Expand Down
168 changes: 168 additions & 0 deletions numba_cuda/numba/cuda/core/codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,168 @@
from abc import abstractmethod, ABCMeta
from numba.misc.llvm_pass_timings import PassTimingsCollection


class CodeLibrary(metaclass=ABCMeta):
"""
An interface for bundling LLVM code together and compiling it.
It is tied to a *codegen* instance (e.g. JITCUDACodegen) that will
determine how the LLVM code is transformed and linked together.
"""

_finalized = False
_object_caching_enabled = False
_disable_inspection = False

def __init__(self, codegen: "Codegen", name: str):
self._codegen = codegen
self._name = name
ptc_name = f"{self.__class__.__name__}({self._name!r})"
self._recorded_timings = PassTimingsCollection(ptc_name)
# Track names of the dynamic globals
self._dynamic_globals = []

@property
def has_dynamic_globals(self):
self._ensure_finalized()
return len(self._dynamic_globals) > 0

@property
def recorded_timings(self):
return self._recorded_timings

@property
def codegen(self):
"""
The codegen object owning this library.
"""
return self._codegen

@property
def name(self):
return self._name

def __repr__(self):
return "<Library %r at 0x%x>" % (self.name, id(self))

def _raise_if_finalized(self):
if self._finalized:
raise RuntimeError(
"operation impossible on finalized object %r" % (self,)
)

def _ensure_finalized(self):
if not self._finalized:
self.finalize()

def create_ir_module(self, name):
"""
Create an LLVM IR module for use by this library.
"""
self._raise_if_finalized()
ir_module = self._codegen._create_empty_module(name)
return ir_module

@abstractmethod
def add_linking_library(self, library):
"""
Add a library for linking into this library, without losing
the original library.
"""

@abstractmethod
def add_ir_module(self, ir_module):
"""
Add an LLVM IR module's contents to this library.
"""

@abstractmethod
def finalize(self):
"""
Finalize the library. After this call, nothing can be added anymore.
Finalization involves various stages of code optimization and
linking.
"""

@abstractmethod
def get_function(self, name):
"""
Return the function named ``name``.
"""

@abstractmethod
def get_llvm_str(self):
"""
Get the human-readable form of the LLVM module.
"""

@abstractmethod
def get_asm_str(self):
"""
Get the human-readable assembly.
"""

#
# Object cache hooks and serialization
#

def enable_object_caching(self):
self._object_caching_enabled = True
self._compiled_object = None
self._compiled = False

def _get_compiled_object(self):
if not self._object_caching_enabled:
raise ValueError("object caching not enabled in %s" % (self,))
if self._compiled_object is None:
raise RuntimeError("no compiled object yet for %s" % (self,))
return self._compiled_object

def _set_compiled_object(self, value):
if not self._object_caching_enabled:
raise ValueError("object caching not enabled in %s" % (self,))
if self._compiled:
raise ValueError("library already compiled: %s" % (self,))
self._compiled_object = value
self._disable_inspection = True


class Codegen(metaclass=ABCMeta):
"""
Base Codegen class. It is expected that subclasses set the class attribute
``_library_class``, indicating the CodeLibrary class for the target.

Subclasses should also initialize:

``self._data_layout``: the data layout for the target.
``self._target_data``: the binding layer ``TargetData`` for the target.
"""

@abstractmethod
def _create_empty_module(self, name):
"""
Create a new empty module suitable for the target.
"""

@abstractmethod
def _add_module(self, module):
"""
Add a module to the execution engine. Ownership of the module is
transferred to the engine.
"""

@property
def target_data(self):
"""
The LLVM "target data" object for this codegen instance.
"""
return self._target_data

def create_library(self, name, **kwargs):
"""
Create a :class:`CodeLibrary` object for use with this codegen
instance.
"""
return self._library_class(self, name, **kwargs)

def unserialize_library(self, serialized):
return self._library_class._unserialize(self, serialized)