yum-mirror/slang
Making it easier to work with shaders
git clone https://git.yummers.dev/yum-mirror/slang
8a8ff3cf3
master
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 ))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
- calling
squareand bindingtorchtensors as arguments for the kernel, and then - launching it using
launchRaw()by specifying CUDA launch arguments toblockSize&gridSize. (Refer to the CUDA documentation for restrictions aroundblockSize)
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 ))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 input_grad )# Note that the derivatives in output_grad are 'consumed'. # i.e. all zeros after the call. 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' )f"X = { x } " )y_pred = MySquareFunc .apply (x )loss = y_pred .sum ()loss .backward ()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
inputGradToPropagateToparameter. - Modified
input[x,y]with a call togetInputElement. - 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 )f"X = { x } " )y = m .square (x )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. 48uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); 49 50// If the thread index is beyond the input size, exit early. 51if (dispatchIdx.x >= input.size(0)) 52return; 53 54output[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{ 111uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); 112 113if (dispatchIdx.x >= input.size(0)) 114return; 115116 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 ( 148input = ( input , input_grad ), output = ( output , output_grad ) 149). launchRaw ( 150blockSize = ( 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 178def forward ( ctx , input ): 179output = torch . zeros_like ( input ) 180 181kernel_with_args = m . square ( input = input , output = output ) 182kernel_with_args . launchRaw ( 183blockSize = ( 32 , 32 , 1 ), 184gridSize = (( input . shape [ 0 ] + 31 ) // 32 , ( input . shape [ 1 ] + 31 ) // 32 , 1 )) 185 186ctx . save_for_backward ( input , output ) 187 188return output 189 190@staticmethod 191def backward ( ctx , grad_output ): 192( input , output ) = ctx . saved_tensors 193 194input_grad = torch . zeros_like ( input ) 195196 # 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# 199kernel_with_args = m . square . bwd ( input = ( input , input_grad ), output = ( output , grad_output )) 200kernel_with_args . launchRaw ( 201blockSize = ( 32 , 32 , 1 ), 202gridSize = (( input . shape [ 0 ] + 31 ) // 32 , ( input . shape [ 1 ] + 31 ) // 32 , 1 )) 203204 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.], 221device='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{ 240var pt = ray.o; 241for (int i = 0; i < N; i++) 242{ 243pt += sdf.eval(pt) * ray.d; 244} 245 246return pt; 247} 248 249float render(Ray ray) 250{ 251// Use N=20 for sphere tracing. 252float3 pt = sphereTrace<20>(ray, sdf); 253return 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{ 265float3 pt = sphereTrace<kMaxSteps>(ray, sdf); 266return 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{ 293int width = input.size(0); 294int height = input.size(1); 295 296// Track the sum of neighboring pixels and the number 297// of pixels currently accumulated. 298int count = 0; 299float sumValue = 0.0; 300 301// Iterate through the surrounding area. 302for (int offsetX = -1; offsetX <= 1; offsetX++) 303{ 304// Skip out of bounds pixels. 305int x = pixelLoc.x + offsetX; 306if (x < 0 || x >= width) continue; 307 308for (int offsetY = -1; offsetY <= 1; offsetY++) 309{ 310int y = pixelLoc.y + offsetY; 311if (y < 0 || y >= height) continue; 312sumValue += input[x, y]; 313count++; 314} 315} 316 317// Compute the average value. 318sumValue /= count; 319 320return 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{ 330uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy; 331int width = input.dim(0); 332int height = input.dim(1); 333if (pixelLoc.x >= width) return; 334if (pixelLoc.y >= height) return; 335 336float outputValueAtPixel = computeOutputPixel(input, pixelLoc) 337 338// Write to output tensor. 339output[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( 353TensorView<float> input, 354TensorView<float> inputGradToPropagateTo, 355uint2 loc) 356{ 357return 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( 371TensorView<float> input, 372TensorView<float> inputGradToPropagateTo, 373uint2 pixelLoc) 374{ 375int width = input.dim(0); 376int height = input.dim(1); 377 378// Track the sum of neighboring pixels and the number 379// of pixels currently accumulated. 380int count = 0; 381float sumValue = 0.0; 382 383// Iterate through the surrounding area. 384for (int offsetX = -1; offsetX <= 1; offsetX++) 385{ 386// Skip out of bounds pixels. 387int x = pixelLoc.x + offsetX; 388if (x < 0 || x >= width) continue; 389 390for (int offsetY = -1; offsetY <= 1; offsetY++) 391{ 392int y = pixelLoc.y + offsetY; 393if (y < 0 || y >= height) continue; 394sumValue += getInputElement(input, inputGradToPropagateTo, uint2(x, y)); 395count++; 396} 397} 398 399// Compute the average value. 400sumValue /= count; 401 402return 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( 416TensorView<float> input, 417TensorView<float> resultGradToPropagateFrom, 418TensorView<float> inputGradToPropagateTo) 419{ 420uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy; 421int width = input.dim(0); 422int height = input.dim(1); 423if (pixelLoc.x >= width) return; 424if (pixelLoc.y >= height) return; 425 426bwd_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( 438TensorView<float> input, 439TensorView<float> inputGradToPropagateTo, 440uint2 loc, 441float derivative) 442{ 443float oldVal; 444inputGradToPropagateTo.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{ 471return x * x; 472} 473 474[CudaKernel] 475void square_kernel(TensorView<float> input, TensorView<float> output) 476{ 477uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx(); 478 479if (globalIdx.x >= input.size(0)) 480return; 481 482float result = compute_square(input[globalIdx.x]); 483 484output[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{ 494var result = TorchTensor<float>.zerosLike(input); 495let blockCount = uint3(1); 496let groupSize = uint3(result.size(0), result.size(1), 1); 497__dispatch_kernel(square_kernel, blockCount, groupSize)(input, result); 498return 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{ 547return 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{ 558uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx(); 559 560if (globalIdx.x >= input.size(0) || globalIdx.y >= input.size(1)) 561return; 562 563DifferentialPair<float> dpInput = diffPair(input[globalIdx.xy]); 564var gradInElem = grad_out[globalIdx.xy]; 565bwd_diff(square)(dpInput, gradInElem); 566grad_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{ 595var grad_propagated = TorchTensor<float>.zerosLike(input); 596let blockCount = uint3(1); 597let groupSize = uint3(input.size(0), input.size(1), 1); 598__dispatch_kernel(square_bwd_kernel, blockCount, groupSize)(input, grad_out, grad_propagated); 599return 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 747 When 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{ 753TensorView<float> vertices; 754TensorView<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{ 797TorchTensor<T> tensors[3]; 798float 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.