C++ Extensions

C++ Extensions#

The C++ extensions in xlnstorch.csrc provide optimized implementations of various operations and utilities that are used in xlnstorch. These are only available if you have built the C++ extensions which is done automatically when you install xlnstorch (if you have a compatible C++ compiler). If the extensions are not built, xlnstorch will fall back to pure Python implementations for those operations, which may be slower but will still work correctly.

To JIT compile and load the C++ extensions at runtime, you can call the xlnstorch.csrc.load_backend() function.

load_backend([build_dir, verbose, enable_cpp])

Try to compile and load the C++ extension if it is not already built and loaded.

You can verify if the C++ extensions are available by checking the xlnstorch.CSRC_AVAILABLE flag. If it is set to True, you can import and use the xlnstorch.csrc module. To switch between the C++ and Python implementations, you can call the xlnstorch.operators.toggle_cpp_implementations() function.

The following functions are available in the xlnstorch.csrc module. They should not be called directly by users, but if necessary, the following docs are provided for reference:

float_to_lns_forward(x, base)

Convert a tensor from floating-point representation to LNS representation.

float_to_lns_backward(grad, base)

Convert the gradient from LNS representation back to floating-point representation.

change_base_forward(x, old_base, new_base)

Change the base of an LNS tensor.

change_base_backward(grad, old_base, new_base)

Change the base of the gradient tensor from LNS representation.

set_default_sbdb_implementation(impl_key)

Sets the default implementation for the SBDB function used in LNS addition.

get_table(ez, sbdb, base)

Gives access to the underlying data pointer of the sbdb Tensor and saves copies of the ez and base tensors.

add_forward(x, y, base)

Perform LNS addition on two tensors internal representations.

sum_forward(x, base, dim, keepdim)

Perform LNS summation over specified dimensions.

matmul_forward(A, B, base)

Perform LNS matrix multiplication.

matmul_backward(grad, A, B, base)

Perform the backward pass for LNS matrix multiplication.

conv1d_forward(x, weight, bias, base, ...)

Perform LNS 1D convolution.

conv1d_backward(grad, x, weight, base, ...)

Perform the backward pass for LNS 1D convolution.

conv2d_forward(x, weight, bias, base, ...)

Perform LNS 2D convolution.

conv2d_backward(grad, x, weight, base, ...)

Perform the backward pass for LNS 2D convolution.

conv3d_forward(x, weight, bias, base, ...)

Perform LNS 3D convolution.

conv3d_backward(grad, x, weight, base, ...)

Perform the backward pass for LNS 3D convolution.

avg_pool1d_forward(x, kernel_size, base, ...)

Perform LNS 1D average pooling.

avg_pool1d_backward(grad, x, kernel_size, ...)

Perform the backward pass for LNS 1D average pooling.