yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Yong HeFix doc. (#6494)8a8ff3cf3

master
32.6 KiB813 linesraw

layout: deprecated permalink: "docs/user-guide/a1-02-slangpy"

Using Slang to Write PyTorch Kernels

Deprecated Feature

Note: This documentation is about slang-torch, an old way to use Slang with Python and PyTorch. Developers who are building new projects should use <a href="https://slangpy.shader-slang.org">SlangPy</a> instead.

If you are a PyTorch user seeking to write complex, high-performance, and automatically differentiated kernel functions using a per-thread programming model, we invite you to try Slang. Slang is a cutting-edge shading language that provides a straightforward way to define kernel functions that run incredibly fast in graphics applications. With the latest addition of automatic differentiation and PyTorch interop features, Slang offers an efficient solution for developing auto-differentiated kernels that run at lightning speed with a strongly typed, per-thread programming model.

One of the primary advantages of a per-thread programming model in kernel programming is the elimination of concerns regarding maintaining masks for branches. When developing a kernel in Slang, you can use all control flow statements, composite data types (structs, arrays, etc.), and function calls without additional effort. Code created with these language constructs can be automatically differentiated by the compiler without any restrictions. Additionally, Slang is a strongly typed language, which ensures that you will never encounter type errors at runtime. Most code errors can be identified as you type thanks to the compiler's coding assistance service, further streamlining the development process.

In addition, using a per-thread programming model also results in more optimized memory usage. When writing a kernel in Slang, most intermediate results do not need to be written out to global memory and then read back, reducing global memory bandwidth consumption and the delay caused by these memory operations. As a result, a Slang kernel can typically run at higher efficiency compared to the traditional bulk-synchronous programming model.

Getting Started with SlangTorch

In this tutorial, we will use a simple example to walk through the steps to use Slang in your PyTorch project.

Installation

slangtorch is available via PyPI, so you can install it simply through

pip install slangtorch

Note that slangtorch requires torch with CUDA support. See the pytorch installation page to find the right version for your platform.

You can check that you have the right installation by running:

python -c "import torch; print(f'cuda: {torch.cuda.is_available()}')"

Writing Slang kernels for slangtorch >= v1.1.5

From v2023.4.0, Slang supports auto-binding features that make it easier than ever to invoke Slang kernels from python, and interoperate seamlessly with pytorch tensors.

Here's a barebones example of a simple squaring kernel written in Slang (square.slang):

[AutoPyBindCUDA]
[CUDAKernel]
void square(TensorView<float> input, TensorView<float> output)
{
    // Get the 'global' index of this thread.
    uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();

    // If the thread index is beyond the input size, exit early.
    if (dispatchIdx.x >= input.size(0))
        return;

    output[dispatchIdx.x] = input[dispatchIdx.x] * input[dispatchIdx.x];
}

This code follows the standard pattern of a typical CUDA kernel function. It takes as input two tensors, input and output. It first obtains the global dispatch index of the current thread and performs range check to make sure we don't read or write out of the bounds of input and output tensors, and then calls square() to compute the per-element result, and store it at the corresponding location in output tensor.

slangtorch works by compiling kernels to CUDA and it identifies the functions to compile by checking for the [CUDAKernel] attribute. The second attribute [AutoPyBindCUDA] allows us to call square directly from python without having to write any host code. If you would like to write the host code yourself for finer control, see the other version of this example here.

You can now simply invoke this kernel from python:

import torch
import slangtorch

m = slangtorch.loadModule('square.slang')

A = torch.randn((1024,), dtype=torch.float).cuda()

output = torch.zeros_like(A).cuda()

# Number of threads launched = blockSize * gridSize
m.square(input=A, output=output).launchRaw(blockSize=(32, 1, 1), gridSize=(64, 1, 1))

print(output)

The python script slangtorch.loadModule("square.slang") returns a scope that contains a handle to the square kernel.

The kernel can be invoked by

  1. calling square and binding torch tensors as arguments for the kernel, and then
  2. launching it using launchRaw() by specifying CUDA launch arguments to blockSize & gridSize. (Refer to the CUDA documentation for restrictions around blockSize)

Note that for semantic clarity reasons, calling a kernel requires the use of keyword arguments with names that are lifted from the .slang implementation.

Invoking derivatives of kernels using slangtorch

The [AutoPyBindCUDA] attribute can also be used on differentiable functions defined in Slang, and will automatically bind the derivatives. To do this, simply add the [Differentiable] attribute.

One key point is that the basic TensorView<T> objects are not differentiable. They can be used as buffers for data that does not require derivatives, or even as buffers for the manual accumulation of derivatives.

Instead, use the DiffTensorView type for when you need differentiable tensors. Currently, DiffTensorView only supports the float dtype variety.

Here's a barebones example of a differentiable version of square:

[AutoPyBindCUDA]
[CUDAKernel]
[Differentiable]
void square(DiffTensorView input, DiffTensorView output)
{
    uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();

    if (dispatchIdx.x >= input.size(0))
        return;
    
    output[dispatchIdx.x] = input[dispatchIdx.x] * input[dispatchIdx.x];
}

Now, slangtorch.loadModule("square.slang") returns a scope with three callable handles square, square.fwd for the forward-mode derivative & square.bwd for the reverse-mode derivative.

You can invoke square() normally to get the same effect as the previous example, or invoke square.fwd() / square.bwd() by binding pairs of tensors to compute the derivatives.

import torch
import slangtorch

m = slangtorch.loadModule('square.slang')

input = torch.tensor((0, 1, 2, 3, 4, 5), dtype=torch.float).cuda()
output = torch.zeros_like(input).cuda()

# Invoke normally
m.square(input=input, output=output).launchRaw(blockSize=(6, 1, 1), gridSize=(1, 1, 1))

print(output)

# Invoke reverse-mode autodiff by first allocating tensors to hold the gradients
input = torch.tensor((0, 1, 2, 3, 4, 5), dtype=torch.float).cuda()
input_grad = torch.zeros_like(input).cuda()

output = torch.zeros_like(input)
# Pass in all 1s as the output derivative for our example
output_grad = torch.ones_like(output) 

m.square.bwd(
    input=(input, input_grad), output=(output, output_grad)
).launchRaw(
    blockSize=(6, 1, 1), gridSize=(1, 1, 1))

# Derivatives get propagated to input_grad
print(input_grad)

# Note that the derivatives in output_grad are 'consumed'.
# i.e. all zeros after the call.
print(output_grad)

slangtorch also binds the forward-mode version of your kernel (propagate derivatives of inputs to the output) which can be invoked the same way using module.square.fwd()

You can refer to this documentation for a detailed reference of Slang's automatic differentiation feature.

Wrapping your kernels as pytorch functions

pytorch offers an easy way to define a custom operation using torch.autograd.Function, and defining the .forward() and .backward() members.

This can be a very helpful way to wrap your Slang kernels as pytorch-compatible operations. Here's an example of the square kernel as a differentiable pytorch function.

import torch
import slangtorch

m = slangtorch.loadModule("square.slang")

class MySquareFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        output = torch.zeros_like(input)

        kernel_with_args = m.square(input=input, output=output)
        kernel_with_args.launchRaw(
            blockSize=(32, 32, 1),
            gridSize=((input.shape[0] + 31) // 32, (input.shape[1] + 31) // 32, 1))

        ctx.save_for_backward(input, output)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        (input, output) = ctx.saved_tensors

        input_grad = torch.zeros_like(input)
        
        # Note: When using DiffTensorView, grad_output gets 'consumed' during the reverse-mode.
        # If grad_output may be reused, consider calling grad_output = grad_output.clone()
        #
        kernel_with_args = m.square.bwd(input=(input, input_grad), output=(output, grad_output))
        kernel_with_args.launchRaw(
            blockSize=(32, 32, 1),
            gridSize=((input.shape[0] + 31) // 32, (input.shape[1] + 31) // 32, 1))
        
        return input_grad

Now we can use the autograd function MySquareFunc in our python script:

x = torch.tensor((3.0, 4.0), requires_grad=True, device='cuda')
print(f"X = {x}")
y_pred = MySquareFunc.apply(x)
loss = y_pred.sum()
loss.backward()
print(f"dX = {x.grad.cpu()}")

Output:

X = tensor([3., 4.],
           device='cuda:0', requires_grad=True)
dX = tensor([6., 8.])

And that's it! slangtorch.loadModule uses JIT compilation to compile your Slang source into CUDA binary. It may take a little longer the first time you execute the script, but the compiled binaries will be cached and as long as the kernel code is not changed, future runs will not rebuild the CUDA kernel.

Because the PyTorch JIT system requires ninja, you need to make sure ninja is installed on your system and is discoverable from the current environment, you also need to have a C++ compiler available on the system. On Windows, this means that Visual Studio need to be installed.

Specializing shaders using slangtorch

slangtorch.loadModule allows specialization parameters to be specified since it might be easier to write shaders with placeholder definitions that can be substituted at load-time. For instance, here's a sphere tracer that uses a compile-time specialization parameter for its maximum number of steps (N):

float sphereTrace<let N:int>(Ray ray, SDF sdf)
{
    var pt = ray.o;
    for (int i = 0; i < N; i++)
    {
        pt += sdf.eval(pt) * ray.d;
    }

    return pt;
}

float render(Ray ray)
{
    // Use N=20 for sphere tracing.
    float3 pt = sphereTrace<20>(ray, sdf);
    return shade(pt, sdf.normal());
}

However, instead of using a fixed 20 steps, the renderer can be configured to use an arbitrary compile-time constant.

// Compile-time constant. Expect "MAX_STEPS" to be set by the loadModule call.
static const uint kMaxSteps = MAX_STEPS;

float render(Ray ray)
{
    float3 pt = sphereTrace<kMaxSteps>(ray, sdf);
    return shade(pt, sdf.normal());
}

Then multiple versions of this shader can be compiled from Python using the defines argument:

import slangtorch

sdfRenderer20Steps = slangtorch.loadModule('sdf.slang', defines={"MAX_STEPS": 20})
sdfRenderer50Steps = slangtorch.loadModule('sdf.slang', defines={"MAX_STEPS": 50})
...

This is often helpful for code re-use, parameter sweeping, comparison/ablation studies, and more, from the convenience of Python.

Back-propagating Derivatives through Complex Access Patterns

In most common scenarios, a kernel function will access input tensors in a complex pattern instead of mapping 1:1 from an input element to an output element, like the square example shown above. When you have a kernel function that access many different elements from the input tensors and use them to compute an output element, the derivatives of each input element can't be represented directly as a function parameter, like the x in square(x).

Consider a 3x3 box filtering kernel that computes for each pixel in a 2D image, the average value of its surrounding 3x3 pixel block. We can write a Slang function that computes the value of an output pixel:

float computeOutputPixel(TensorView<float> input, uint2 pixelLoc)
{
    int width = input.size(0);
    int height = input.size(1);

    // Track the sum of neighboring pixels and the number
    // of pixels currently accumulated.
    int count = 0;
    float sumValue = 0.0;

    // Iterate through the surrounding area.
    for (int offsetX = -1; offsetX <= 1; offsetX++)
    {
        // Skip out of bounds pixels.
        int x = pixelLoc.x + offsetX;
        if (x < 0 || x >= width) continue;

        for (int offsetY = -1; offsetY <= 1; offsetY++)
        {
            int y = pixelLoc.y + offsetY;
            if (y < 0 || y >= height) continue;
            sumValue += input[x, y];
            count++;
        }
    }

    // Compute the average value.
    sumValue /= count;

    return sumValue;
}

We can define our kernel function to compute the entire output image by calling computeOutputPixel:

[CudaKernel]
void boxFilter_fwd(TensorView<float> input, TensorView<float> output)
{
    uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy;
    int width = input.dim(0);
    int height = input.dim(1);
    if (pixelLoc.x >= width) return;
    if (pixelLoc.y >= height) return;

    float outputValueAtPixel = computeOutputPixel(input, pixelLoc)

    // Write to output tensor.
    output[pixelLoc] = outputValueAtPixel;
}

How do we define the backward derivative propagation kernel? Note that in this example, there isn't a function like square that we can just mark as [Differentiable] and call bwd_diff(square) to get back the derivative of an input parameter.

In this example, the input comes from multiple elements in a tensor. How do we propagate the derivatives to those input elements?

The solution is to wrap tensor access with a custom function:

float getInputElement(
    TensorView<float> input,
    TensorView<float> inputGradToPropagateTo,
    uint2 loc)
{
    return input[loc];
}

Note that the getInputElement function simply returns input[loc] and is not using the inputGradToPropagateTo parameter. That is intended. The inputGradToPropagateTo parameter is used to hold the backward propagated derivatives of each input element, and is reserved for later use.

Now we can replace all direct accesses to input with a call to getInputElement. The computeOutputPixel can be implemented as following:

[Differentiable]
float computeOutputPixel(
    TensorView<float> input,
    TensorView<float> inputGradToPropagateTo,
    uint2 pixelLoc)
{
    int width = input.dim(0);
    int height = input.dim(1);

    // Track the sum of neighboring pixels and the number
    // of pixels currently accumulated.
    int count = 0;
    float sumValue = 0.0;

    // Iterate through the surrounding area.
    for (int offsetX = -1; offsetX <= 1; offsetX++)
    {
        // Skip out of bounds pixels.
        int x = pixelLoc.x + offsetX;
        if (x < 0 || x >= width) continue;

        for (int offsetY = -1; offsetY <= 1; offsetY++)
        {
            int y = pixelLoc.y + offsetY;
            if (y < 0 || y >= height) continue;
            sumValue += getInputElement(input, inputGradToPropagateTo, uint2(x, y));
            count++;
        }
    }

    // Compute the average value.
    sumValue /= count;

    return sumValue;
}

The main changes compared to our original version of computeOutputPixel are:

  • Added a inputGradToPropagateTo parameter.
  • Modified input[x,y] with a call to getInputElement.
  • Added a [Differentiable] attribute to the function.

With that, we can define our backward kernel function:

[CudaKernel]
void boxFilter_bwd(
    TensorView<float> input,
    TensorView<float> resultGradToPropagateFrom,
    TensorView<float> inputGradToPropagateTo)
{
    uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy;
    int width = input.dim(0);
    int height = input.dim(1);
    if (pixelLoc.x >= width) return;
    if (pixelLoc.y >= height) return;

    bwd_diff(computeOutputPixel)(input, inputGradToPropagateTo, pixelLoc);
}

The kernel function simply calls bwd_diff(computeOutputPixel) without taking any return values from the call and without writing to any elements in the final inputGradToPropagateTo tensor. But when exactly does the propagated output get written to the output gradient tensor (inputGradToPropagateTo)?

And that logic is defined in our final piece of code:

[BackwardDerivativeOf(getInputElement)]
void getInputElement_bwd(
    TensorView<float> input,
    TensorView<float> inputGradToPropagateTo,
    uint2 loc,
    float derivative)
{
    float oldVal;
    inputGradToPropagateTo.InterlockedAdd(loc, derivative, oldVal);
}

Here, we are providing a custom defined backward propagation function for getInputElement. In this function, we simply add derivative to the element in inputGradToPropagateTo tensor.

When we call bwd_diff(computeOutputPixel) in boxFilter_bwd, the Slang compiler will automatically differentiate all operations and function calls in computeOutputPixel. By wrapping the tensor element access with getInputElement and by providing a custom backward propagation function of getInputElement, we are effectively telling the compiler what to do when a derivative propagates to an input tensor element. Inside the body of getInputElement_bwd, we define what to do then: atomically adds the derivative propagated to the input element in the inputGradToPropagateTo tensor. Therefore, after running boxFilter_bwd, the inputGradToPropagateTo tensor will contain all the back propagated derivative values.

Again, to understand all the details of the automatic differentiation system, please refer to the Automatic Differentiation chapter for a detailed explanation.

Manually binding kernels

[AutoPyBindCUDA] works for most use cases, but in certain situations, it may be necessary to write the host function by hand. The host function can also be written in Slang, and slangtorch handles its compilation to C++.

Here's the same square example from before:

// square.slang
float compute_square(float x)
{
    return x * x;
}

[CudaKernel]
void square_kernel(TensorView<float> input, TensorView<float> output)
{
    uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();

    if (globalIdx.x >= input.size(0))
        return;

    float result = compute_square(input[globalIdx.x]);

    output[globalIdx.x] = result;
}

To manually invoke this kernel, we then need to write a CPU(host) function that defines how this kernel is dispatched. This can be defined in the same Slang file:

[TorchEntryPoint]
TorchTensor<float> square(TorchTensor<float> input)
{
    var result = TorchTensor<float>.zerosLike(input);
    let blockCount = uint3(1);
    let groupSize = uint3(result.size(0), result.size(1), 1);
    __dispatch_kernel(square_kernel, blockCount, groupSize)(input, result);
    return result;
}

Here, we mark the function with the [TorchEntryPoint] attribute, so it will be compiled to C++ and exported as a python callable. Since this is a host function, we can perform tensor allocations. For instance, square() calls TorchTensor<float>.zerosLike to allocate a 2D-tensor that has the same size as the input. zerosLike returns a TorchTensor<float> object that represents a CPU handle of a PyTorch tensor.

Then we launch square_kernel with the __dispatch_kernel syntax. Note that we can directly pass TorchTensor<float> arguments to a TensorView<float> parameter and the compiler will automatically convert the type and obtain a view into the tensor that can be accessed by the GPU kernel function.

Calling a [TorchEntryPoint] function from Python

You can use the following code to call square from Python:

import torch
import slangtorch

m = slangtorch.loadModule("square.slang")

x = torch.randn(2,2)
print(f"X = {x}")
y = m.square(x)
print(f"Y = {y.cpu()}")

Result output:

X = tensor([[ 0.1407,  0.6594],
        [-0.8978, -1.7230]])
Y = tensor([[0.0198, 0.4349],
        [0.8060, 2.9688]])

Manual binding for kernel derivatives

The above example demonstrates how to write a simple kernel function in Slang and call it from Python. Another major benefit of using Slang is that the Slang compiler support generating backward derivative propagation functions automatically.

In the following section, we walk through how to use Slang to generate a backward propagation function for square, and expose it to PyTorch as an autograd function.

First we need to tell Slang compiler that we need the square function to be considered a differentiable function, so Slang compiler can generate a backward derivative propagation function for it:

[Differentiable]
float square(float x)
{
    return x * x;
}

This is done by simply adding a [Differentiable] attribute to our square function.

With that, we can now define square_bwd_kernel that performs backward propagation as:

[CudaKernel]
void square_bwd_kernel(TensorView<float> input, TensorView<float> grad_out, TensorView<float> grad_propagated)
{
    uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();

    if (globalIdx.x >= input.size(0) || globalIdx.y >= input.size(1))
        return;

    DifferentialPair<float> dpInput = diffPair(input[globalIdx.xy]);
    var gradInElem = grad_out[globalIdx.xy];
    bwd_diff(square)(dpInput, gradInElem);
    grad_propagated[globalIdx.xy] = dpInput.d;
}

Note that the function follows the same structure of square_fwd_kernel, with the only difference being that instead of calling into square to compute the forward value for each tensor element, we are calling bwd_diff(square) that represents the automatically generated backward propagation function of square. bwd_diff(square) will have the following signature:

void bwd_diff_square(inout DifferentialPair<float> dpInput, float dOut);

Where the first parameter, dpInput represents a pair of original and derivative value for input, and the second parameter, dOut, represents the initial derivative with regard to some latent variable that we wish to back-prop through. The resulting derivative will be stored in dpInput.d. For example:

// construct a pair where the primal value is 3, and derivative value is 0.
var dp = diffPair(3.0);
bwd_diff(square)(dp, 1.0);
// dp.d is now 6.0

Similar to square_fwd, we can define the host side function square_bwd as:

[TorchEntryPoint]
TorchTensor<float> square_bwd(TorchTensor<float> input, TorchTensor<float> grad_out)
{
    var grad_propagated = TorchTensor<float>.zerosLike(input);
    let blockCount = uint3(1);
    let groupSize = uint3(input.size(0), input.size(1), 1);
    __dispatch_kernel(square_bwd_kernel, blockCount, groupSize)(input, grad_out, grad_propagated);
    return grad_propagated;
}

Builtin Library Support for PyTorch Interop

As shown in previous tutorial, Slang has defined the TorchTensor<T> and TensorView<T> type for interop with PyTorch tensors. The TorchTensor<T> represents the CPU view of a tensor and provides methods to allocate a new tensor object. The TensorView<T> represents the GPU view of a tensor and provides accessors to read write tensor data.

Following is a list of built-in methods and attributes for PyTorch interop.

TorchTensor methods

static TorchTensor<T> TorchTensor<T>.alloc(uint x, uint y, ...)

Allocates a new PyTorch tensor with the given dimensions. If T is a vector type, the length of the vector is implicitly included as the last dimension. For example, TorchTensor<float3>.alloc(4, 4) allocates a 3D tensor of size (4,4,3).

static TorchTensor<T> TorchTensor<T>.emptyLike(TorchTensor<T> other)

Allocates a new PyTorch tensor that has the same dimensions as other without initializing it.

static TorchTensor<T> TorchTensor<T>.zerosLike(TorchTensor<T> other)

Allocates a new PyTorch tensor that has the same dimensions as other and initialize it to zero.

uint TorchTensor<T>.dims()

Returns the tensor's dimension count.

uint TorchTensor<T>.size(int dim)

Returns the tensor's size (in number of elements) at dim.

uint TorchTensor<T>.stride(int dim)

Returns the tensor's stride (in bytes) at dim.

TensorView methods

TensorView<T>.operator[uint x, uint y, ...]

Provide an accessor to data content in a tensor.

TensorView<T>.operator[vector<uint, N> index]

Provide an accessor to data content in a tensor, indexed by a uint vector. tensor[uint3(1,2,3)] is equivalent to tensor[1,2,3].

uint TensorView<T>.dims()

Returns the tensor's dimension count.

uint TensorView<T>.size(int dim)

Returns the tensor's size (in number of elements) at dim.

uint TensorView<T>.stride(int dim)

Returns the tensor's stride (in bytes) at dim.

void TensorView<T>.fillZero()

Fills the tensor with zeros. Modifies the tensor in-place.

void TensorView<T>.fillValue(T value)

Fills the tensor with the specified value, modifies the tensor in-place.

T* TensorView<T>.data_ptr_at(vector<uint, N> index)

Returns a pointer to the element at index.

void TensorView<T>.InterlockedAdd(vector<uint, N> index, T val, out T oldVal)

Atomically add val to element at index.

void TensorView<T>.InterlockedMin(vector<uint, N> index, T val, out T oldVal)

Atomically computes the min of val and the element at index. Available for 32 and 64 bit integer types only.

void TensorView<T>.InterlockedMax(vector<uint, N> index, T val, out T oldVal)

Atomically computes the max of val and the element at index. Available for 32 and 64 bit integer types only.

void TensorView<T>.InterlockedAnd(vector<uint, N> index, T val, out T oldVal)

Atomically computes the bitwise and of val and the element at index. Available for 32 and 64 bit integer types only.

void TensorView<T>.InterlockedOr(vector<uint, N> index, T val, out T oldVal)

Atomically computes the bitwise or of val and the element at index. Available for 32 and 64 bit integer types only.

void TensorView<T>.InterlockedXor(vector<uint, N> index, T val, out T oldVal)

Atomically computes the bitwise xor of val and the element at index. Available for 32 and 64 bit integer types only.

void TensorView<T>.InterlockedExchange(vector<uint, N> index, T val, out T oldVal)

Atomically swaps val into the element at index. Available for float and 32/64 bit integer types only.

void TensorView<T>.InterlockedCompareExchange(vector<uint, N> index, T compare, T val)

Atomically swaps val into the element at index if the element equals to compare. Available for float and 32/64 bit integer types only.

DiffTensorView methods

DiffTensorView.operator[uint x, uint y, ...]

Provide an accessor to data content in a tensor. This method is differentiable, and has the same semantics as using a .load() to get data, and .store() to set data.

DiffTensorView.operator[vector<uint, N> index]

Provide an accessor to data content in a tensor, indexed by a uint vector.tensor[uint3(1,2,3)] is equivalent to tensor[1,2,3]. This method is differentiable, and has the same semantics as using a .load() to get data, and .store() to set data.

float DiffTensorView.load(vector<uint, N> index)

Loads the 32-bit floating point data at the specified multi-dimensional index. This method is differentiable, and in reverse-mode will perform an atomic-add.

void DiffTensorView.store(vector<uint, N> index, float val)

Stores the 32-bit floating point value val at the specified multi-dimensional index. This method is differentiable, and in reverse-mode will perform an atomic exchange to retrieve the derivative and replace with 0.

float DiffTensorView.loadOnce(vector<uint, N> index)

Loads the 32-bit floating point data at the specified multi-dimensional index. This method is differentiable, and uses a simple store for the reverse-mode for faster gradient aggregation, but loadOnce must be used at most once per index. loadOnce is ideal for situations where each thread loads data from a unique index, but will cause incorrect gradients when an index may be accessed multiple times.

void DiffTensorView.storeOnce(vector<uint, N> index, float val)

Stores the 32-bit floating point value val at the specified multi-dimensional index. This method is differentiable, and uses a simple load for the reverse-mode for faster gradient loading, but storeOnce must be used at most once per index. loadOnce is ideal for situations where each thread stores data to a unique index, but will cause incorrect gradient propagation when an index may be accessed multiple times.

uint DiffTensorView.size(int dim)

Returns the underlying primal tensor's size (in number of elements) at dim.

uint DiffTensorView.dims()

Returns the underlying primal tensor's dimension count.

uint DiffTensorView.stride(uint dim)

Returns the stride of the underlying primal tensor's dim dimension

CUDA Support Functions

cudaThreadIdx()

Returns the threadIdx variable in CUDA.

cudaBlockIdx()

Returns the blockIdx variable in CUDA.

cudaBlockDim()

Returns the blockDim variable in CUDA.

syncTorchCudaStream()

Waits for all pending CUDA kernel executions to complete on host.

Attributes for PyTorch Interop

[CudaKernel] attribute

Marks a function as a CUDA kernel (maps to a __global__ function)

[TorchEntryPoint] attribute

Marks a function for export to Python. Functions marked with [TorchEntryPoint] will be accessible from a loaded module returned by slangtorch.loadModule.

[CudaDeviceExport] attribute

Marks a function as a CUDA device function, and ensures the compiler to include it in the generated CUDA source.

[AutoPyBindCUDA] attribute

Marks a cuda kernel for automatic binding generation so that it may be invoked from python without having to hand-code the torch entry point. The marked function must also be marked with [CudaKernel]. If the marked function is also marked with [Differentiable], this will also generate bindings for the derivative methods.

Restriction: methods marked with [AutoPyBindCUDA] will not operate

Type Marshalling Between Slang and Python

Python-CUDA type marshalling for functions using [AutoPyBindCUDA]

When using auto-binding, aggregate types like structs are converted to Python namedtuples and are made available when using slangtorch.loadModule.

// mesh.slang
struct Mesh
{
    TensorView<float> vertices;
    TensorView<int> indices;
};

[AutoPyBindCUDA]
[CUDAKernel]
void processMesh(Mesh mesh)
{
    /* ... */ 
}

Here, since Mesh is being used by renderMesh, the loaded module will provide Mesh as a python namedtuple with named fields. While using the namedtuple is the best way to use structured arguments, they can also be passed as a python dict or tuple

m = slangtorch.loadModule('mesh.slang')

vertices = torch.tensor()
indices = torch.tensor()

# use namedtuple to provide structured input.
mesh = m.Mesh(vertices=vertices, indices=indices)
m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))

# use dict to provide input.
mesh = {'vertices': vertices, 'indices':indices}
m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))

# use tuple to provide input (warning: user responsible for right order)
mesh = (vertices, indices)
m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))

Python-CUDA type marshalling for functions using [TorchEntryPoint]

The return types and parameters types of an exported [TorchEntryPoint] function can be a basic type (e.g. float, int etc.), a vector type (e.g. float3), a TorchTensor<T> type, an array type, or a struct type.

When you use struct or array types in the function signature, it will be exposed as a Python tuple. For example,

struct MyReturnType
{
    TorchTensor<T> tensors[3];
    float v;
}

[TorchEntryPoint]
MyReturnType myFunc()
{
    ...
}

Calling myFunc from python will result in a python tuple in the form of

[[tensor, tensor, tensor], float]

The same transform rules apply to parameter types.

1---
2layout: deprecated
3permalink: "docs/user-guide/a1-02-slangpy"
4---
5
6Using Slang to Write PyTorch Kernels
7=========================================================
8
9> #### Deprecated Feature
10> Note: This documentation is about `slang-torch`, an old way to use Slang with Python and PyTorch.
11> Developers who are building new projects should use <a href="https://slangpy.shader-slang.org">SlangPy</a> instead.
12
13If you are a PyTorch user seeking to write complex, high-performance, and automatically differentiated kernel functions using a per-thread programming model, we invite you to try Slang. Slang is a cutting-edge shading language that provides a straightforward way to define kernel functions that run incredibly fast in graphics applications. With the latest addition of automatic differentiation and PyTorch interop features, Slang offers an efficient solution for developing auto-differentiated kernels that run at lightning speed with a strongly typed, per-thread programming model.
14
15One of the primary advantages of a per-thread programming model in kernel programming is the elimination of concerns regarding maintaining masks for branches. When developing a kernel in Slang, you can use all control flow statements, composite data types (structs, arrays, etc.), and function calls without additional effort. Code created with these language constructs can be automatically differentiated by the compiler without any restrictions. Additionally, Slang is a strongly typed language, which ensures that you will never encounter type errors at runtime. Most code errors can be identified as you type thanks to the [compiler's coding assistance service](https://marketplace.visualstudio.com/items?itemName=shader-slang.slang-language-extension), further streamlining the development process.
16
17In addition, using a per-thread programming model also results in more optimized memory usage. When writing a kernel in Slang, most intermediate results do not need to be written out to global memory and then read back, reducing global memory bandwidth consumption and the delay caused by these memory operations. As a result, a Slang kernel can typically run at higher efficiency compared to the traditional bulk-synchronous programming model.
18
19## Getting Started with SlangTorch
20
21In this tutorial, we will use a simple example to walk through the steps to use Slang in your PyTorch project.
22
23### Installation
24`slangtorch` is available via PyPI, so you can install it simply through
25```sh
26pip install slangtorch
27```
28
29Note that `slangtorch` requires `torch` with CUDA support. See the [pytorch](https://pytorch.org/) installation page to find the right version for your platform.
30
31You can check that you have the right installation by running: 
32```sh
33python -c "import torch; print(f'cuda: {torch.cuda.is_available()}')"
34```
35
36### Writing Slang kernels for `slangtorch` >= **v1.1.5**
37
38From **v2023.4.0**, Slang supports auto-binding features that make it easier than ever to invoke Slang kernels from python, and interoperate seamlessly with `pytorch` tensors.
39
40Here's a barebones example of a simple squaring kernel written in Slang (`square.slang`):
41
42```csharp
43[AutoPyBindCUDA]
44[CUDAKernel]
45void square(TensorView<float> input, TensorView<float> output)
46{
47    // Get the 'global' index of this thread.
48    uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();
49
50    // If the thread index is beyond the input size, exit early.
51    if (dispatchIdx.x >= input.size(0))
52        return;
53
54    output[dispatchIdx.x] = input[dispatchIdx.x] * input[dispatchIdx.x];
55}
56
57```
58
59This code follows the standard pattern of a typical CUDA kernel function. It takes as input
60two tensors, `input` and `output`. 
61It first obtains the global dispatch index of the current thread and performs range check to make sure we don't read or write out
62of the bounds of input and output tensors, and then calls `square()` to compute the per-element result, and
63store it at the corresponding location in `output` tensor.
64
65
66`slangtorch` works by compiling kernels to CUDA and it identifies the functions to compile by checking for the `[CUDAKernel]` attribute.
67The second attribute `[AutoPyBindCUDA]` allows us to call `square` directly from python without having to write any host code. If you would like to write the host code yourself for finer control, see the other version of this example [here](#manually-binding-kernels).
68
69You can now simply invoke this kernel from python:
70
71```python
72import torch
73import slangtorch
74
75m = slangtorch.loadModule('square.slang')
76
77A = torch.randn((1024,), dtype=torch.float).cuda()
78
79output = torch.zeros_like(A).cuda()
80
81# Number of threads launched = blockSize * gridSize
82m.square(input=A, output=output).launchRaw(blockSize=(32, 1, 1), gridSize=(64, 1, 1))
83
84print(output)
85```
86
87The python script `slangtorch.loadModule("square.slang")` returns a scope that contains a handle to the `square` kernel.
88
89The kernel can be invoked by 
901. calling `square` and binding `torch` tensors as arguments for the kernel, and then
912. launching it using `launchRaw()` by specifying CUDA launch arguments to `blockSize` & `gridSize`. (Refer to the [CUDA documentation](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications) for restrictions around `blockSize`)
92
93Note that for semantic clarity reasons, calling a kernel requires the use of keyword arguments with names that are lifted from the `.slang` implementation.
94
95### Invoking derivatives of kernels using slangtorch
96
97The `[AutoPyBindCUDA]` attribute can also be used on differentiable functions defined in Slang, and will automatically bind the derivatives. To do this, simply add the `[Differentiable]` attribute.
98
99One key point is that the basic `TensorView<T>` objects are not differentiable. They can be used as buffers for data that does not require derivatives, or even as buffers for the manual accumulation of derivatives.
100
101Instead, use the `DiffTensorView` type for when you need differentiable tensors. Currently, `DiffTensorView` only supports the `float` dtype variety.
102
103Here's a barebones example of a differentiable version of `square`:
104
105```csharp
106[AutoPyBindCUDA]
107[CUDAKernel]
108[Differentiable]
109void square(DiffTensorView input, DiffTensorView output)
110{
111    uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();
112
113    if (dispatchIdx.x >= input.size(0))
114        return;
115    
116    output[dispatchIdx.x] = input[dispatchIdx.x] * input[dispatchIdx.x];
117}
118```
119
120Now, `slangtorch.loadModule("square.slang")` returns a scope with three callable handles `square`, `square.fwd` for the forward-mode derivative & `square.bwd` for the reverse-mode derivative.
121
122You can invoke `square()` normally to get the same effect as the previous example, or invoke `square.fwd()` / `square.bwd()` by binding pairs of tensors to compute the derivatives.
123
124
125```python
126import torch
127import slangtorch
128
129m = slangtorch.loadModule('square.slang')
130
131input = torch.tensor((0, 1, 2, 3, 4, 5), dtype=torch.float).cuda()
132output = torch.zeros_like(input).cuda()
133
134# Invoke normally
135m.square(input=input, output=output).launchRaw(blockSize=(6, 1, 1), gridSize=(1, 1, 1))
136
137print(output)
138
139# Invoke reverse-mode autodiff by first allocating tensors to hold the gradients
140input = torch.tensor((0, 1, 2, 3, 4, 5), dtype=torch.float).cuda()
141input_grad = torch.zeros_like(input).cuda()
142
143output = torch.zeros_like(input)
144# Pass in all 1s as the output derivative for our example
145output_grad = torch.ones_like(output) 
146
147m.square.bwd(
148    input=(input, input_grad), output=(output, output_grad)
149).launchRaw(
150    blockSize=(6, 1, 1), gridSize=(1, 1, 1))
151
152# Derivatives get propagated to input_grad
153print(input_grad)
154
155# Note that the derivatives in output_grad are 'consumed'.
156# i.e. all zeros after the call.
157print(output_grad)
158```
159
160`slangtorch` also binds the forward-mode version of your kernel (propagate derivatives of inputs to the output) which can be invoked the same way using `module.square.fwd()`
161
162You can refer to [this documentation](autodiff) for a detailed reference of Slang's automatic differentiation feature.
163
164### Wrapping your kernels as pytorch functions
165
166`pytorch` offers an easy way to define a custom operation using `torch.autograd.Function`, and defining the `.forward()` and `.backward()` members.
167
168This can be a very helpful way to wrap your Slang kernels as pytorch-compatible operations. Here's an example of the `square` kernel as a differentiable pytorch function.
169
170```python
171import torch
172import slangtorch
173
174m = slangtorch.loadModule("square.slang")
175
176class MySquareFunc(torch.autograd.Function):
177    @staticmethod
178    def forward(ctx, input):
179        output = torch.zeros_like(input)
180
181        kernel_with_args = m.square(input=input, output=output)
182        kernel_with_args.launchRaw(
183            blockSize=(32, 32, 1),
184            gridSize=((input.shape[0] + 31) // 32, (input.shape[1] + 31) // 32, 1))
185
186        ctx.save_for_backward(input, output)
187
188        return output
189
190    @staticmethod
191    def backward(ctx, grad_output):
192        (input, output) = ctx.saved_tensors
193
194        input_grad = torch.zeros_like(input)
195        
196        # Note: When using DiffTensorView, grad_output gets 'consumed' during the reverse-mode.
197        # If grad_output may be reused, consider calling grad_output = grad_output.clone()
198        #
199        kernel_with_args = m.square.bwd(input=(input, input_grad), output=(output, grad_output))
200        kernel_with_args.launchRaw(
201            blockSize=(32, 32, 1),
202            gridSize=((input.shape[0] + 31) // 32, (input.shape[1] + 31) // 32, 1))
203        
204        return input_grad
205```
206
207Now we can use the autograd function `MySquareFunc` in our python script:
208
209```python
210x = torch.tensor((3.0, 4.0), requires_grad=True, device='cuda')
211print(f"X = {x}")
212y_pred = MySquareFunc.apply(x)
213loss = y_pred.sum()
214loss.backward()
215print(f"dX = {x.grad.cpu()}")
216```
217
218Output:
219```
220X = tensor([3., 4.],
221           device='cuda:0', requires_grad=True)
222dX = tensor([6., 8.])
223```
224
225And that's it! `slangtorch.loadModule` uses JIT compilation to compile your Slang source into CUDA binary.
226It may take a little longer the first time you execute the script, but the compiled binaries will be cached and as long as the kernel code is not changed, future runs will not rebuild the CUDA kernel.
227
228Because the PyTorch JIT system requires `ninja`, you need to make sure `ninja` is installed on your system
229and is discoverable from the current environment, you also need to have a C++ compiler available on the system.
230On Windows, this means that Visual Studio need to be installed.
231
232## Specializing shaders using slangtorch
233
234`slangtorch.loadModule` allows specialization parameters to be specified since it might be easier to write shaders with placeholder definitions that can be substituted at load-time.
235For instance, here's a sphere tracer that uses a _compile-time_ specialization parameter for its maximum number of steps (`N`):
236
237```csharp
238float sphereTrace<let N:int>(Ray ray, SDF sdf)
239{
240    var pt = ray.o;
241    for (int i = 0; i < N; i++)
242    {
243        pt += sdf.eval(pt) * ray.d;
244    }
245
246    return pt;
247}
248
249float render(Ray ray)
250{
251    // Use N=20 for sphere tracing.
252    float3 pt = sphereTrace<20>(ray, sdf);
253    return shade(pt, sdf.normal());
254}
255```
256
257However, instead of using a fixed `20` steps, the renderer can be configured to use an arbitrary compile-time constant.
258
259```csharp
260// Compile-time constant. Expect "MAX_STEPS" to be set by the loadModule call.
261static const uint kMaxSteps = MAX_STEPS;
262
263float render(Ray ray)
264{
265    float3 pt = sphereTrace<kMaxSteps>(ray, sdf);
266    return shade(pt, sdf.normal());
267}
268```
269
270Then multiple versions of this shader can be compiled from Python using the `defines` argument:
271```python
272import slangtorch
273
274sdfRenderer20Steps = slangtorch.loadModule('sdf.slang', defines={"MAX_STEPS": 20})
275sdfRenderer50Steps = slangtorch.loadModule('sdf.slang', defines={"MAX_STEPS": 50})
276...
277```
278
279This is often helpful for code re-use, parameter sweeping, comparison/ablation studies, and more, from the convenience of Python.
280
281## Back-propagating Derivatives through Complex Access Patterns
282
283In most common scenarios, a kernel function will access input tensors in a complex pattern instead of mapping
2841:1 from an input element to an output element, like the `square` example shown above. When you have a kernel
285function that access many different elements from the input tensors and use them to compute an output element,
286the derivatives of each input element can't be represented directly as a function parameter, like the `x` in `square(x)`.
287
288Consider a 3x3 box filtering kernel that computes for each pixel in a 2D image, the average value of its 
289surrounding 3x3 pixel block. We can write a Slang function that computes the value of an output pixel:
290```csharp
291float computeOutputPixel(TensorView<float> input, uint2 pixelLoc)
292{
293    int width = input.size(0);
294    int height = input.size(1);
295
296    // Track the sum of neighboring pixels and the number
297    // of pixels currently accumulated.
298    int count = 0;
299    float sumValue = 0.0;
300
301    // Iterate through the surrounding area.
302    for (int offsetX = -1; offsetX <= 1; offsetX++)
303    {
304        // Skip out of bounds pixels.
305        int x = pixelLoc.x + offsetX;
306        if (x < 0 || x >= width) continue;
307
308        for (int offsetY = -1; offsetY <= 1; offsetY++)
309        {
310            int y = pixelLoc.y + offsetY;
311            if (y < 0 || y >= height) continue;
312            sumValue += input[x, y];
313            count++;
314        }
315    }
316
317    // Compute the average value.
318    sumValue /= count;
319
320    return sumValue;
321}
322```
323
324We can define our kernel function to compute the entire output image by calling `computeOutputPixel`:
325
326```csharp
327[CudaKernel]
328void boxFilter_fwd(TensorView<float> input, TensorView<float> output)
329{
330    uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy;
331    int width = input.dim(0);
332    int height = input.dim(1);
333    if (pixelLoc.x >= width) return;
334    if (pixelLoc.y >= height) return;
335
336    float outputValueAtPixel = computeOutputPixel(input, pixelLoc)
337
338    // Write to output tensor.
339    output[pixelLoc] = outputValueAtPixel;
340}
341```
342
343How do we define the backward derivative propagation kernel? Note that in this example, there
344isn't a function like `square` that we can just mark as `[Differentiable]` and
345call `bwd_diff(square)` to get back the derivative of an input parameter.
346
347In this example, the input comes from multiple elements in a tensor. How do we propagate the
348derivatives to those input elements?
349
350The solution is to wrap tensor access with a custom function:
351```csharp
352float getInputElement(
353    TensorView<float> input,
354    TensorView<float> inputGradToPropagateTo,
355    uint2 loc)
356{
357    return input[loc];
358}
359```
360
361Note that the `getInputElement` function simply returns `input[loc]` and is not using the
362`inputGradToPropagateTo` parameter. That is intended. The `inputGradToPropagateTo` parameter
363is used to hold the backward propagated derivatives of each input element, and is reserved for later use.
364
365Now we can replace all direct accesses to `input` with a call to `getInputElement`. The
366`computeOutputPixel` can be implemented as following:
367
368```csharp
369[Differentiable]
370float computeOutputPixel(
371    TensorView<float> input,
372    TensorView<float> inputGradToPropagateTo,
373    uint2 pixelLoc)
374{
375    int width = input.dim(0);
376    int height = input.dim(1);
377
378    // Track the sum of neighboring pixels and the number
379    // of pixels currently accumulated.
380    int count = 0;
381    float sumValue = 0.0;
382
383    // Iterate through the surrounding area.
384    for (int offsetX = -1; offsetX <= 1; offsetX++)
385    {
386        // Skip out of bounds pixels.
387        int x = pixelLoc.x + offsetX;
388        if (x < 0 || x >= width) continue;
389
390        for (int offsetY = -1; offsetY <= 1; offsetY++)
391        {
392            int y = pixelLoc.y + offsetY;
393            if (y < 0 || y >= height) continue;
394            sumValue += getInputElement(input, inputGradToPropagateTo, uint2(x, y));
395            count++;
396        }
397    }
398
399    // Compute the average value.
400    sumValue /= count;
401
402    return sumValue;
403}
404```
405
406The main changes compared to our original version of `computeOutputPixel` are:
407- Added a `inputGradToPropagateTo` parameter.
408- Modified `input[x,y]` with a call to `getInputElement`.
409- Added a `[Differentiable]` attribute to the function.
410
411With that, we can define our backward kernel function:
412
413```csharp
414[CudaKernel]
415void boxFilter_bwd(
416    TensorView<float> input,
417    TensorView<float> resultGradToPropagateFrom,
418    TensorView<float> inputGradToPropagateTo)
419{
420    uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy;
421    int width = input.dim(0);
422    int height = input.dim(1);
423    if (pixelLoc.x >= width) return;
424    if (pixelLoc.y >= height) return;
425
426    bwd_diff(computeOutputPixel)(input, inputGradToPropagateTo, pixelLoc);
427}
428```
429
430The kernel function simply calls `bwd_diff(computeOutputPixel)` without taking any return values from the call
431and without writing to any elements in the final `inputGradToPropagateTo` tensor. But when exactly does the propagated
432output get written to the output gradient tensor (`inputGradToPropagateTo`)?
433
434And that logic is defined in our final piece of code:
435```csharp
436[BackwardDerivativeOf(getInputElement)]
437void getInputElement_bwd(
438    TensorView<float> input,
439    TensorView<float> inputGradToPropagateTo,
440    uint2 loc,
441    float derivative)
442{
443    float oldVal;
444    inputGradToPropagateTo.InterlockedAdd(loc, derivative, oldVal);
445}
446```
447
448Here, we are providing a custom defined backward propagation function for `getInputElement`.
449In this function, we simply add `derivative` to the element in `inputGradToPropagateTo` tensor.
450
451When we call `bwd_diff(computeOutputPixel)` in `boxFilter_bwd`, the Slang compiler will automatically
452differentiate all operations and function calls in `computeOutputPixel`. By wrapping the tensor element access
453with `getInputElement` and by providing a custom backward propagation function of `getInputElement`, we are effectively
454telling the compiler what to do when a derivative propagates to an input tensor element. Inside the body
455of `getInputElement_bwd`, we define what to do then: atomically adds the derivative propagated to the input element
456in the `inputGradToPropagateTo` tensor. Therefore, after running `boxFilter_bwd`, the `inputGradToPropagateTo` tensor will contain all the
457back propagated derivative values.
458
459Again, to understand all the details of the automatic differentiation system, please refer to the 
460[Automatic Differentiation](autodiff) chapter for a detailed explanation.
461
462## Manually binding kernels
463`[AutoPyBindCUDA]` works for most use cases, but in certain situations, it may be necessary to write the *host* function by hand. The host function can also be written in Slang, and `slangtorch` handles its compilation to C++.
464
465Here's the same `square` example from before:
466
467```csharp
468// square.slang
469float compute_square(float x)
470{
471    return x * x;
472}
473
474[CudaKernel]
475void square_kernel(TensorView<float> input, TensorView<float> output)
476{
477    uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();
478
479    if (globalIdx.x >= input.size(0))
480        return;
481
482    float result = compute_square(input[globalIdx.x]);
483
484    output[globalIdx.x] = result;
485}
486```
487
488To manually invoke this kernel, we then need to write a CPU(host) function that defines how this kernel is dispatched. This can be defined in the same Slang file:
489
490```csharp
491[TorchEntryPoint]
492TorchTensor<float> square(TorchTensor<float> input)
493{
494    var result = TorchTensor<float>.zerosLike(input);
495    let blockCount = uint3(1);
496    let groupSize = uint3(result.size(0), result.size(1), 1);
497    __dispatch_kernel(square_kernel, blockCount, groupSize)(input, result);
498    return result;
499}
500```
501
502Here, we mark the function with the `[TorchEntryPoint]` attribute, so it will be compiled to C++ and exported as a python callable. 
503Since this is a host function, we can perform tensor allocations. For instance, `square()` calls `TorchTensor<float>.zerosLike` to allocate a 2D-tensor that has the same size as the input.
504`zerosLike` returns a `TorchTensor<float>` object that represents a CPU handle of a PyTorch tensor.
505
506Then we launch `square_kernel` with the `__dispatch_kernel` syntax. Note that we can directly pass
507`TorchTensor<float>` arguments to a `TensorView<float>` parameter and the compiler will automatically convert the type and obtain a view into the tensor that can be accessed by the GPU kernel function.
508
509### Calling a `[TorchEntryPoint]` function from Python
510
511You can use the following code to call `square` from Python:
512
513```python
514import torch
515import slangtorch
516
517m = slangtorch.loadModule("square.slang")
518
519x = torch.randn(2,2)
520print(f"X = {x}")
521y = m.square(x)
522print(f"Y = {y.cpu()}")
523```
524
525Result output:
526```
527X = tensor([[ 0.1407,  0.6594],
528        [-0.8978, -1.7230]])
529Y = tensor([[0.0198, 0.4349],
530        [0.8060, 2.9688]])
531```
532
533### Manual binding for kernel derivatives
534
535The above example demonstrates how to write a simple kernel function in Slang and call it from Python.
536Another major benefit of using Slang is that the Slang compiler support generating backward derivative
537propagation functions automatically.
538
539In the following section, we walk through how to use Slang to generate a backward propagation function
540for `square`, and expose it to PyTorch as an autograd function.
541
542First we need to tell Slang compiler that we need the `square` function to be considered a differentiable function, so Slang compiler can generate a backward derivative propagation function for it:
543```csharp
544[Differentiable]
545float square(float x)
546{
547    return x * x;
548}
549```
550This is done by simply adding a `[Differentiable]` attribute to our `square` function.
551
552With that, we can now define `square_bwd_kernel` that performs backward propagation as:
553
554```csharp
555[CudaKernel]
556void square_bwd_kernel(TensorView<float> input, TensorView<float> grad_out, TensorView<float> grad_propagated)
557{
558    uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();
559
560    if (globalIdx.x >= input.size(0) || globalIdx.y >= input.size(1))
561        return;
562
563    DifferentialPair<float> dpInput = diffPair(input[globalIdx.xy]);
564    var gradInElem = grad_out[globalIdx.xy];
565    bwd_diff(square)(dpInput, gradInElem);
566    grad_propagated[globalIdx.xy] = dpInput.d;
567}
568```
569
570Note that the function follows the same structure of `square_fwd_kernel`, with the only difference being that
571instead of calling into `square` to compute the forward value for each tensor element, we are calling `bwd_diff(square)`
572that represents the automatically generated backward propagation function of `square`.
573`bwd_diff(square)` will have the following signature:
574```csharp
575void bwd_diff_square(inout DifferentialPair<float> dpInput, float dOut);
576```
577
578Where the first parameter, `dpInput` represents a pair of original and derivative value for `input`, and the second parameter,
579`dOut`, represents the initial derivative with regard to some latent variable that we wish to back-prop through. The resulting
580derivative will be stored in `dpInput.d`. For example:
581
582```csharp
583// construct a pair where the primal value is 3, and derivative value is 0.
584var dp = diffPair(3.0);
585bwd_diff(square)(dp, 1.0);
586// dp.d is now 6.0
587```
588
589Similar to `square_fwd`, we can define the host side function `square_bwd` as:
590
591```csharp
592[TorchEntryPoint]
593TorchTensor<float> square_bwd(TorchTensor<float> input, TorchTensor<float> grad_out)
594{
595    var grad_propagated = TorchTensor<float>.zerosLike(input);
596    let blockCount = uint3(1);
597    let groupSize = uint3(input.size(0), input.size(1), 1);
598    __dispatch_kernel(square_bwd_kernel, blockCount, groupSize)(input, grad_out, grad_propagated);
599    return grad_propagated;
600}
601```
602
603## Builtin Library Support for PyTorch Interop
604
605As shown in previous tutorial, Slang has defined the `TorchTensor<T>` and `TensorView<T>` type for interop with PyTorch
606tensors. The `TorchTensor<T>` represents the CPU view of a tensor and provides methods to allocate a new tensor object.
607The `TensorView<T>` represents the GPU view of a tensor and provides accessors to read write tensor data.
608
609Following is a list of built-in methods and attributes for PyTorch interop.
610
611### `TorchTensor` methods
612
613#### `static TorchTensor<T> TorchTensor<T>.alloc(uint x, uint y, ...)`
614Allocates a new PyTorch tensor with the given dimensions. If `T` is a vector type, the length of the vector is implicitly included as the last dimension.
615For example, `TorchTensor<float3>.alloc(4, 4)` allocates a 3D tensor of size `(4,4,3)`.
616
617#### `static TorchTensor<T> TorchTensor<T>.emptyLike(TorchTensor<T> other)`
618Allocates a new PyTorch tensor that has the same dimensions as `other` without initializing it.
619
620#### `static TorchTensor<T> TorchTensor<T>.zerosLike(TorchTensor<T> other)`
621Allocates a new PyTorch tensor that has the same dimensions as `other` and initialize it to zero.
622
623#### `uint TorchTensor<T>.dims()`
624Returns the tensor's dimension count.
625
626#### `uint TorchTensor<T>.size(int dim)`
627Returns the tensor's size (in number of elements) at `dim`.
628
629#### `uint TorchTensor<T>.stride(int dim)`
630Returns the tensor's stride (in bytes) at `dim`.
631
632### `TensorView` methods
633
634#### `TensorView<T>.operator[uint x, uint y, ...]`
635Provide an accessor to data content in a tensor.
636
637#### `TensorView<T>.operator[vector<uint, N> index]`
638Provide an accessor to data content in a tensor, indexed by a uint vector.
639`tensor[uint3(1,2,3)]` is equivalent to `tensor[1,2,3]`.
640
641#### `uint TensorView<T>.dims()`
642Returns the tensor's dimension count.
643
644#### `uint TensorView<T>.size(int dim)`
645Returns the tensor's size (in number of elements) at `dim`.
646
647#### `uint TensorView<T>.stride(int dim)`
648Returns the tensor's stride (in bytes) at `dim`.
649
650#### `void TensorView<T>.fillZero()`
651Fills the tensor with zeros. Modifies the tensor in-place.
652
653#### `void TensorView<T>.fillValue(T value)`
654Fills the tensor with the specified value, modifies the tensor in-place.
655
656#### `T* TensorView<T>.data_ptr_at(vector<uint, N> index)`
657Returns a pointer to the element at `index`.
658
659#### `void TensorView<T>.InterlockedAdd(vector<uint, N> index, T val, out T oldVal)`
660Atomically add `val` to element at `index`. 
661
662#### `void TensorView<T>.InterlockedMin(vector<uint, N> index, T val, out T oldVal)`
663Atomically computes the min of `val` and the element at `index`. Available for 32 and 64 bit integer types only.
664
665#### `void TensorView<T>.InterlockedMax(vector<uint, N> index, T val, out T oldVal)`
666Atomically computes the max of `val` and the element at `index`. Available for 32 and 64 bit integer types only.
667
668#### `void TensorView<T>.InterlockedAnd(vector<uint, N> index, T val, out T oldVal)`
669Atomically computes the bitwise and of `val` and the element at `index`. Available for 32 and 64 bit integer types only.
670
671#### `void TensorView<T>.InterlockedOr(vector<uint, N> index, T val, out T oldVal)`
672Atomically computes the bitwise or  of `val` and the element at `index`. Available for 32 and 64 bit integer types only.
673
674#### `void TensorView<T>.InterlockedXor(vector<uint, N> index, T val, out T oldVal)`
675Atomically computes the bitwise xor  of `val` and the element at `index`. Available for 32 and 64 bit integer types only.
676
677#### `void TensorView<T>.InterlockedExchange(vector<uint, N> index, T val, out T oldVal)`
678Atomically swaps `val` into the element at `index`. Available for `float` and 32/64 bit integer types only.
679
680#### `void TensorView<T>.InterlockedCompareExchange(vector<uint, N> index, T compare, T val)`
681Atomically swaps `val` into the element at `index` if the element equals to `compare`. Available for `float` and 32/64 bit integer types only.
682
683### `DiffTensorView` methods
684
685#### `DiffTensorView.operator[uint x, uint y, ...]`
686Provide an accessor to data content in a tensor. This method is **differentiable**, and has the same semantics as using a `.load()` to get data, and `.store()` to set data.
687
688#### `DiffTensorView.operator[vector<uint, N> index]`
689Provide an accessor to data content in a tensor, indexed by a uint vector.`tensor[uint3(1,2,3)]` is equivalent to `tensor[1,2,3]`. This method is **differentiable**, and has the same semantics as using a `.load()` to get data, and `.store()` to set data.
690
691#### `float DiffTensorView.load(vector<uint, N> index)`
692Loads the 32-bit floating point data at the specified multi-dimensional `index`. This method is **differentiable**, and in reverse-mode will perform an atomic-add.
693
694#### `void DiffTensorView.store(vector<uint, N> index, float val)`
695Stores the 32-bit floating point value `val` at the specified multi-dimensional `index`. This method is **differentiable**, and in reverse-mode will perform an *atomic exchange* to retrieve the derivative and replace with 0.
696
697#### `float DiffTensorView.loadOnce(vector<uint, N> index)`
698Loads the 32-bit floating point data at the specified multi-dimensional `index`. This method is **differentiable**, and uses a simple `store` for the reverse-mode for faster gradient aggregation, but `loadOnce` **must** be used at most once per index. `loadOnce` is ideal for situations where each thread loads data from a unique index, but will cause incorrect gradients when an index may be accessed multiple times.
699
700#### `void DiffTensorView.storeOnce(vector<uint, N> index, float val)`
701Stores the 32-bit floating point value `val` at the specified multi-dimensional `index`. This method is **differentiable**, and uses a simple `load` for the reverse-mode for faster gradient loading, but `storeOnce` **must** be used at most once per index. `loadOnce` is ideal for situations where each thread stores data to a unique index, but will cause incorrect gradient propagation when an index may be accessed multiple times.
702
703#### `uint DiffTensorView.size(int dim)`
704Returns the underlying primal tensor's size (in number of elements) at `dim`.
705
706#### `uint DiffTensorView.dims()`
707Returns the underlying primal tensor's dimension count.
708
709#### `uint DiffTensorView.stride(uint dim)`
710Returns the stride of the underlying primal tensor's `dim` dimension
711
712### CUDA Support Functions
713
714#### `cudaThreadIdx()`
715Returns the `threadIdx` variable in CUDA.
716
717#### `cudaBlockIdx()`
718Returns the `blockIdx` variable in CUDA.
719
720#### `cudaBlockDim()`
721Returns the `blockDim` variable in CUDA.
722
723#### `syncTorchCudaStream()`
724Waits for all pending CUDA kernel executions to complete on host.
725
726### Attributes for PyTorch Interop
727
728#### `[CudaKernel]` attribute
729Marks a function as a CUDA kernel (maps to a `__global__` function)
730
731#### `[TorchEntryPoint]` attribute
732Marks a function for export to Python. Functions marked with `[TorchEntryPoint]` will be accessible from a loaded module returned by `slangtorch.loadModule`.
733
734#### `[CudaDeviceExport]` attribute
735Marks a function as a CUDA device function, and ensures the compiler to include it in the generated CUDA source.
736
737#### `[AutoPyBindCUDA]` attribute
738Marks a cuda kernel for automatic binding generation so that it may be invoked from python without having to hand-code the torch entry point. The marked function **must** also be marked with `[CudaKernel]`. If the marked function is also marked with `[Differentiable]`, this will also generate bindings for the derivative methods.
739
740Restriction: methods marked with `[AutoPyBindCUDA]` will not operate 
741
742## Type Marshalling Between Slang and Python
743
744
745### Python-CUDA type marshalling for functions using `[AutoPyBindCUDA]` 
746
747When using auto-binding, aggregate types like structs are converted to Python `namedtuples` and are made available when using `slangtorch.loadModule`. 
748
749```csharp
750// mesh.slang
751struct Mesh
752{
753    TensorView<float> vertices;
754    TensorView<int> indices;
755};
756
757[AutoPyBindCUDA]
758[CUDAKernel]
759void processMesh(Mesh mesh)
760{
761    /* ... */ 
762}
763```
764
765Here, since `Mesh` is being used by `renderMesh`, the loaded module will provide `Mesh` as a python `namedtuple` with named fields.
766While using the `namedtuple` is the best way to use structured arguments, they can also be passed as a python `dict` or `tuple`
767
768```python
769m = slangtorch.loadModule('mesh.slang')
770
771vertices = torch.tensor()
772indices = torch.tensor()
773
774# use namedtuple to provide structured input.
775mesh = m.Mesh(vertices=vertices, indices=indices)
776m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))
777
778# use dict to provide input.
779mesh = {'vertices': vertices, 'indices':indices}
780m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))
781
782# use tuple to provide input (warning: user responsible for right order)
783mesh = (vertices, indices)
784m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))
785```
786
787
788### Python-CUDA type marshalling for functions using `[TorchEntryPoint]`
789
790The return types and parameters types of an exported `[TorchEntryPoint]` function can be a basic type (e.g. `float`, `int` etc.), a vector type (e.g. `float3`), a `TorchTensor<T>` type, an array type, or a struct type.
791
792When you use struct or array types in the function signature, it will be exposed as a Python tuple.
793For example,
794```csharp
795struct MyReturnType
796{
797    TorchTensor<T> tensors[3];
798    float v;
799}
800
801[TorchEntryPoint]
802MyReturnType myFunc()
803{
804    ...
805}
806```
807
808Calling `myFunc` from python will result in a python tuple in the form of
809```
810[[tensor, tensor, tensor], float]
811```
812
813The same transform rules apply to parameter types.