Core APIs

Core APIs#

Tensor#

class nkipy.core.tensor.NKIPyTensorRef(backend_tensor, name: str = None)[source]#

NKIPy Tensor Reference

This class uses mixins to provide arithmetic operations and implements the numpy array protocol to intercept operations during tracing.

The tensor is backend-agnostic - it holds a reference to a backend tensor (e.g., HLOTensor for HLO backend) and delegates operations to the ops module which dispatches to the appropriate backend implementation.

__init__(backend_tensor, name: str = None)[source]#

Initialize tensor reference

Parameters:

  • backend_tensor – the target backend tensor, e.g., HLOTensor for HLO backend

  • name – Optional name for the tensor

backend_tensor#
property name#
property shape: Tuple[int, ...]#
property dtype#
property ndim: int#
astype(dtype)[source]#

Convert tensor to a different dtype

Compile APIs#

Compiler wrappers to lower NKIPy kernels

class nkipy.core.compile.CompilationTarget(*values)[source]#
TRN1 = 'trn1'#
TRN2 = 'trn2'#
DEFAULT = 'default'#
nkipy.core.compile.get_platform_target() → str[source]#
class nkipy.core.compile.CompilationConfig(pipeline: tuple[str] = ('compile', 'SaveTemps'), target: CompilationTarget = CompilationTarget.DEFAULT, additional_args: str = '', neff_name: str = 'file.neff')[source]#

Configuration for compilation process

pipeline: tuple[str] = ('compile', 'SaveTemps')#
target: CompilationTarget = 'default'#
additional_args: str = ''#
neff_name: str = 'file.neff'#
__init__(pipeline: tuple[str] = ('compile', 'SaveTemps'), target: CompilationTarget = CompilationTarget.DEFAULT, additional_args: str = '', neff_name: str = 'file.neff') → None#
class nkipy.core.compile.CompilationResult(work_dir: Path, neff_path: Path)[source]#

Results from compilation process

__init__(work_dir: Path, neff_path: Path)[source]#
save_artifacts(output_dir: Path) → None[source]#

Save all compilation artifacts by copying the entire working directory

class nkipy.core.compile.Compiler(config: CompilationConfig)[source]#

Handles compilation of traced kernels

__init__(config: CompilationConfig)[source]#
compile(ir, work_dir: Path, output_file: str, use_neuronx_cc_python_interface: bool = False) → Path[source]#

Run compilation in specified directory

Parameters:

  • ir – The IR to compile

  • work_dir – Directory to compile in

  • output_file – Name of the output file to check for (“file.neff” or “nki.py”)

Returns:

Path to the output file

compile_in_directory(ir, output_file: str, output_dir: str | None = None, save_artifacts: bool = False, use_neuronx_cc_python_interface: bool = False) → Path[source]#

Compile in either a temporary directory or the specified output directory

Parameters:

  • ir – The IR to compile

  • output_file – Name of the output file to check for

  • output_dir – Directory to save outputs

  • save_artifacts – If True, saves all compilation artifacts

Returns:

Path to the output file

nkipy.core.compile.get_ir(traced_kernel)[source]#
nkipy.core.compile.compile_to_neff(trace_kernel, output_dir: str, neff_name: str = 'file.neff', target: CompilationTarget = CompilationTarget.DEFAULT, additional_compiler_args: str = '', save_artifacts: bool = False, use_neuronx_cc_python_interface: bool = False) → str[source]#

Compile traced kernel to NEFF file

Parameters:

  • trace_kernel – The kernel to compile

  • output_dir – Directory to save outputs

  • target – Target platform for compilation

  • save_artifacts – If True, saves all compilation artifacts

Returns:

Path to the generated NEFF file as str

nkipy.core.compile.lower_to_nki(trace_kernel, output_dir: str | None = None, target: CompilationTarget = CompilationTarget.DEFAULT, additional_compiler_args: str = '', save_artifacts: bool = False) → str[source]#

Lower traced kernel to NKI representation

Parameters:

  • trace_kernel – The kernel to lower

  • output_dir – Directory to save outputs

  • target – Target platform for compilation

  • save_artifacts – If True, saves all compilation artifacts

Returns:

The NKI representation as string

Contents