## 1. Project Overview & Quickstart (shader-slang/slang) ## File: README.md Slang ===== Slang is a shading language that makes it easier to build and maintain large shader codebases in a modular and extensible fashion, while also maintaining the highest possible performance on modern GPUs and graphics APIs. Slang is based on years of collaboration between researchers at NVIDIA, Carnegie Mellon University, Stanford, MIT, UCSD and the University of Washington. Why Slang? --------------- The Slang shading language is designed to enable real-time graphics developers to work with large-scale, high-performance shader code. ### Write Shaders Once, Run Anywhere The Slang compiler can generate code for a wide variety of targets: D3D12, Vulkan, Metal, D3D11, CUDA, and even generate code to run on a CPU. For textual targets, such as Metal Shading Language (MSL) and CUDA, Slang produces readable code that preserves original identifier names, as well as the type and call structure, making it easier to debug. ### Access the Latest GPU Features Slang code is highly portable, but can still leverage unique platform capabilities, including the latest features in Direct3D and Vulkan. For example, developers can make full use of [pointers](https://shader-slang.com/slang/user-guide/convenience-features.html#pointers-limited) when generating SPIR-V. Slang's [capability system](https://shader-slang.com/slang/user-guide/capabilities.html) helps applications manage feature set differences across target platforms by ensuring code only uses available features during the type-checking step, before generating final code. Additionally, Slang provides [flexible interop](https://shader-slang.com/slang/user-guide/a1-04-interop.html) features to enable directly embedding target code or SPIR-V into generated shaders. ### Leverage Neural Graphics with Automatic Differentiation Slang can [automatically generate both forward and backward derivative propagation code](https://shader-slang.com/slang/user-guide/autodiff.html) for complex functions that involve arbitrary control flow and dynamic dispatch. This allows existing rendering codebases to easily become differentiable, or for Slang to serve as the kernel language in a PyTorch-driven machine learning framework via [`slangtorch`](https://shader-slang.com/slang/user-guide/a1-02-slangpy.html). ### Scalable Software Development with Modules Slang provides a [module system](https://shader-slang.com/slang/user-guide/modules.html) that enables logical organization of code for separate compilation. Slang modules can be independently compiled offline to a custom IR (with optional obfuscation) and then linked at runtime to generate code in formats such as DXIL or SPIR-V. ### Code Specialization that Works with Modules Slang supports [generics and interfaces](https://shader-slang.com/slang/user-guide/interfaces-generics.html) (a.k.a. type traits/protocols), allowing for clear expression of shader specialization without the need for preprocessor techniques or string-pasting. Unlike C++ templates, Slang's generics are pre-checked and don't produce cascading error messages that are difficult to diagnose. The same generic shader can be specialized for a variety of different types to produce specialized code ahead of time, or on the fly, entirely under application control. ### Easy On-ramp for HLSL and GLSL Codebases Slang's syntax is similar to HLSL, and most existing HLSL code can be compiled with the Slang compiler out-of-the-box, or with just minor modifications. This allows existing shader codebases to immediately benefit from Slang without requiring a complete rewrite or port. Slang provides a compatibility module that enables the use of most GLSL intrinsic functions and GLSL's parameter binding syntax. ### Comprehensive Tooling Support Slang comes with full support of IntelliSense editing features in Visual Studio Code and Visual Studio through the Language Server Protocol. Full debugging capabilities are also available through RenderDoc and SPIR-V based tools. Getting Started --------------- The fastest way to get started using Slang in your own development is to use a pre-built binary package, available through GitHub [releases](https://github.com/shader-slang/slang/releases). Slang binaries are also included in the [Vulkan SDK](https://vulkan.lunarg.com/sdk/home) since version 1.3.296.0. There are packages built for x86_64 and aarch64 Windows, Linux and macOS. Each binary release includes the command-line `slangc` compiler, a shared library for the compiler, and the `slang.h` header. See the user-guide for info on using the `slangc` command-line tool: [Slang Command Line Usage]( https://shader-slang.com/slang/user-guide/compiling.html#command-line-compilation-with-slangc). If you want to try out the Slang language without installing anything, a fast and simple way is to use the [Slang Playground](https://shader-slang.com/slang-playground). The playground allows you to compile Slang code to a variety of targets, and even run some simple shaders directly within the browser. The playground loads Slang compiler to your browser and runs all compilation locally. No data will be sent to any servers. If you would like to build Slang from source, please consult the [build instructions](docs/building.md). Documentation ------------- The Slang project provides a variety of different [documentation](docs/), but most users would be well served starting with the [User's Guide](https://shader-slang.github.io/slang/user-guide/). For developers writing Slang code, the [Slang Core Module Reference](https://shader-slang.com/stdlib-reference/) provides detailed documentation on Slang's built-in types and functions. We also provide a few [examples](examples/) of how to integrate Slang into a rendering application. These examples use a graphics layer that we include with Slang called "GFX" which is an abstraction library of various graphics APIs (D3D11, D2D12, OpenGL, Vulkan, CUDA, and the CPU) to support cross-platform applications using GPU graphics and compute capabilities. GFX is being deprecated in favor of [slang-rhi](https://github.com/shader-slang/slang-rhi). Additionally, we recommend checking out [Vulkan Mini Examples](https://github.com/nvpro-samples/vk_mini_samples/) for more examples of using Slang's language features available on Vulkan, such as pointers and the ray tracing intrinsics. Contributing ------------ If you'd like to contribute to the project, we are excited to have your input. The following guidelines should be observed by contributors: * Please follow the contributor [Code of Conduct](CODE_OF_CONDUCT.md). * Bugs reports and feature requests should go through the GitHub issue tracker * Changes should ideally come in as small pull requests on top of `master`, coming from your own personal fork of the project * Large features that will involve multiple contributors or a long development time should be discussed in issues, and broken down into smaller pieces that can be implemented and checked in in stages [Contribution guide](CONTRIBUTING.md) describes the workflow for contributors at more detail. Limitations and Support ----------------------- ### Platform support The Slang compiler and libraries can be built on the following platforms: | Windows | Linux | MacOS | WebAssembly | |:---------:|:---------:|:---------:|:------------:| | supported | supported | supported | experimental | Both `x86_64` and `aarch64` architectures are supported on Windows, Linux and MacOS platforms. ### Target support Slang can compile shader code to the following targets: | Target | Status | Output Formats | |:-----------:|:-------------------------------------------------------------------------------------:|:----------------------------------------------------------------:| | Direct3D 11 | [supported](https://shader-slang.com/slang/user-guide/targets.html#direct3d-11) | HLSL | | Direct3D 12 | [supported](https://shader-slang.com/slang/user-guide/targets.html#direct3d-12) | HLSL | | Vulkan | [supported](https://shader-slang.com/slang/user-guide/targets.html#vulkan) | SPIRV, GLSL | | Metal | [experimental*](https://shader-slang.com/slang/user-guide/targets.html#metal) | Metal Shading Language | | WebGPU | experimental** | WGSL | | CUDA | [supported](https://shader-slang.com/slang/user-guide/targets.html#cuda-and-optix) | C++ (compute only) | | Optix | [experimental](https://shader-slang.com/slang/user-guide/targets.html#cuda-and-optix) | C++ (WIP) | | CPU | [experimental](https://shader-slang.com/slang/user-guide/targets.html#cpu-compute) | C++ (kernel), C++ (host), standalone executable, dynamic library | > *Slang currently supports generating vertex, fragment, compute, task and mesh > shaders for Metal. > **WGSL support is still work in-progress. For greater detail, see the [Supported Compilation Targets](https://shader-slang.com/slang/user-guide/targets.html) section of the [User Guide](https://shader-slang.github.io/slang/user-guide/) The Slang project has been used for production applications and large shader codebases, but it is still under active development. Support is currently focused on the platforms (Windows, Linux) and target APIs (Direct3D 12, Vulkan) where Slang is used most heavily. Users who are looking for support on other platforms or APIs should coordinate with the development team via the issue tracker to make sure that their use cases can be supported. License ------- The Slang code itself is under the Apache 2.0 with LLVM Exception license (see [LICENSE](LICENSE)). Builds of the core Slang tools depend on the following projects, either automatically or optionally, which may have their own licenses: * [`ankerl::unordered_dense::{map, set}`](https://github.com/martinus/unordered_dense) (MIT) * [`fast_float`](https://github.com/fastfloat/fast_float) (Apache 2.0 / MIT / Boost) * [`glslang`](https://github.com/KhronosGroup/glslang) (BSD) * [`lz4`](https://github.com/lz4/lz4) (BSD) * [`miniz`](https://github.com/richgel999/miniz) (MIT) * [`spirv-headers`](https://github.com/KhronosGroup/SPIRV-Headers) (Modified MIT) * [`spirv-tools`](https://github.com/KhronosGroup/SPIRV-Tools) (Apache 2.0) Slang releases may include [LLVM](https://github.com/llvm/llvm-project) under the license: * [`llvm`](https://llvm.org/docs/DeveloperPolicy.html#new-llvm-project-license-framework) (Apache 2.0 License with LLVM exceptions) Some of the tests and example programs that build with Slang use the following projects, which may have their own licenses: * [`glm`](https://github.com/g-truc/glm) (MIT) * `stb_image` and `stb_image_write` from the [`stb`](https://github.com/nothings/stb) collection of single-file libraries (Public Domain) * [`tinyobjloader`](https://github.com/tinyobjloader/tinyobjloader) (MIT) --- ## File: docs/deprecated/a1-02-slangpy.md --- layout: deprecated permalink: "docs/user-guide/a1-02-slangpy" --- Using Slang to Write PyTorch Kernels ========================================================= > #### Note > This documentation is about `slang-torch`, a way to use Slang with Python and PyTorch. > For new projects, we recommend exploring [SlangPy](https://slangpy.shader-slang.org) as an alternative. > We plan to deprecate `slang-torch` in favor of SlangPy in the near future, and we will communicate any plans in advance. 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](https://marketplace.visualstudio.com/items?itemName=shader-slang.slang-language-extension), 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 ```sh pip install slangtorch ``` Note that `slangtorch` requires `torch` with CUDA support. See the [pytorch](https://pytorch.org/) installation page to find the right version for your platform. You can check that you have the right installation by running: ```sh 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`): ```csharp [AutoPyBindCUDA] [CUDAKernel] void square(TensorView input, TensorView 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](#manually-binding-kernels). You can now simply invoke this kernel from python: ```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](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications) 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` 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`: ```csharp [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. ```python 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](autodiff) 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. ```python 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: ```python 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`): ```csharp float sphereTrace(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. ```csharp // 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(ray, sdf); return shade(pt, sdf.normal()); } ``` Then multiple versions of this shader can be compiled from Python using the `defines` argument: ```python 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: ```csharp float computeOutputPixel(TensorView 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`: ```csharp [CudaKernel] void boxFilter_fwd(TensorView input, TensorView 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: ```csharp float getInputElement( TensorView input, TensorView 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: ```csharp [Differentiable] float computeOutputPixel( TensorView input, TensorView 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: ```csharp [CudaKernel] void boxFilter_bwd( TensorView input, TensorView resultGradToPropagateFrom, TensorView 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: ```csharp [BackwardDerivativeOf(getInputElement)] void getInputElement_bwd( TensorView input, TensorView 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](autodiff) 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: ```csharp // square.slang float compute_square(float x) { return x * x; } [CudaKernel] void square_kernel(TensorView input, TensorView 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: ```csharp [TorchEntryPoint] TorchTensor square(TorchTensor input) { var result = TorchTensor.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.zerosLike` to allocate a 2D-tensor that has the same size as the input. `zerosLike` returns a `TorchTensor` 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` arguments to a `TensorView` 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: ```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: ```csharp [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: ```csharp [CudaKernel] void square_bwd_kernel(TensorView input, TensorView grad_out, TensorView grad_propagated) { uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx(); if (globalIdx.x >= input.size(0) || globalIdx.y >= input.size(1)) return; DifferentialPair 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: ```csharp void bwd_diff_square(inout DifferentialPair 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: ```csharp // 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: ```csharp [TorchEntryPoint] TorchTensor square_bwd(TorchTensor input, TorchTensor grad_out) { var grad_propagated = TorchTensor.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` and `TensorView` type for interop with PyTorch tensors. The `TorchTensor` represents the CPU view of a tensor and provides methods to allocate a new tensor object. The `TensorView` 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 TorchTensor.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.alloc(4, 4)` allocates a 3D tensor of size `(4,4,3)`. #### `static TorchTensor TorchTensor.emptyLike(TorchTensor other)` Allocates a new PyTorch tensor that has the same dimensions as `other` without initializing it. #### `static TorchTensor TorchTensor.zerosLike(TorchTensor other)` Allocates a new PyTorch tensor that has the same dimensions as `other` and initialize it to zero. #### `uint TorchTensor.dims()` Returns the tensor's dimension count. #### `uint TorchTensor.size(int dim)` Returns the tensor's size (in number of elements) at `dim`. #### `uint TorchTensor.stride(int dim)` Returns the tensor's stride (in bytes) at `dim`. ### `TensorView` methods #### `TensorView.operator[uint x, uint y, ...]` Provide an accessor to data content in a tensor. #### `TensorView.operator[vector 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.dims()` Returns the tensor's dimension count. #### `uint TensorView.size(int dim)` Returns the tensor's size (in number of elements) at `dim`. #### `uint TensorView.stride(int dim)` Returns the tensor's stride (in bytes) at `dim`. #### `void TensorView.fillZero()` Fills the tensor with zeros. Modifies the tensor in-place. #### `void TensorView.fillValue(T value)` Fills the tensor with the specified value, modifies the tensor in-place. #### `T* TensorView.data_ptr_at(vector index)` Returns a pointer to the element at `index`. #### `void TensorView.InterlockedAdd(vector index, T val, out T oldVal)` Atomically add `val` to element at `index`. #### `void TensorView.InterlockedMin(vector 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.InterlockedMax(vector 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.InterlockedAnd(vector 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.InterlockedOr(vector 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.InterlockedXor(vector 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.InterlockedExchange(vector 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.InterlockedCompareExchange(vector 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 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 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 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 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 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`. ```csharp // mesh.slang struct Mesh { TensorView vertices; TensorView 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` ```python 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` 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, ```csharp struct MyReturnType { TorchTensor 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. --- ## File: docs/design/autodiff/basics.md This documentation is intended for Slang contributors and is written from a compiler engineering point of view. For Slang users, see the user-guide at this link: [https://shader-slang.com/slang/user-guide/autodiff.html](https://shader-slang.com/slang/user-guide/autodiff.html) ## What is Automatic Differentiation? Before diving into the design of the automatic differentiation (for brevity, we will call it 'auto-diff') passes, it is important to understand the end goal of what auto-diff tries to achieve. The over-arching goal of Slang's auto-diff is to enable the user to compute derivatives of a given shader program or function's output w.r.t its input parameters. This critical compiler feature enables users to quickly use their shaders with gradient-based parameter optimization algorithms, which forms the backbone of modern machine learning systems. It enables users to train and deploy graphics systems that contain ML primitives (like multi-layer perceptron's or MLPs) or use their shader programs as differentiable primitives within larger ML pipelines. ### More Resources Here are some links to resources that talk more about differentiable programming from a more mathematical perspective: 1. UCSD CSE 291 (Spring 2024): https://cseweb.ucsd.edu/~tzli/cse291/sp2024/ 2. UW CSE 5990 (Winter 2024): https://sites.google.com/cs.washington.edu/cse-599o-dppl ## Definition of Derivatives This section is based off of these slides: https://cseweb.ucsd.edu/~tzli/cse291/sp2024/lectures/03_forward_mode.pdf. Here, we establish the mathematical definition of derivatives, starting with a simple 1D case (function with a single input and output), and extending to the general case of functions mapping multiple inputs to multiple outputs. To avoid confusion, we will denote mathematical functions using LaTeX italic script ($f$, $g$, etc..) and programs that compute these functions with markdown code (`f`, `g`, etc..) ### Derivatives of scalar (1D) functions Consider the simplest case: a smooth scalar mathematical function that maps a real number to another real number: $$f : \mathbb{R} \to \mathbb{R}$$ There are several definitions for a derivative, but we will use the definition that a derivative is the *closest linear approximation* of the output function at a given input location. Concretely, given a specific input $x$, we can create a linear approximation of the function $f$ around $x$ as follows: $$ f(x + dx) \approx f(x) + Df(x) \cdot dx $$ This can also be understood as a geometric 'tangent' to the function at $x$. $Df(x)$ is the slope of $f$ at $x$, i.e. $\frac{\partial f}{\partial x}$, and $dx$ is the perturbation away from $x$. Our approximation is linear as a function of the perturbation $dx$. Note that no matter how non-linear or complex the underlying function $f(x)$ is, the approximation is always linear (this property becomes very important later). ### Forward-mode derivative functions Now consider a concrete program `f` that computes some function. ```C // Computes square of x float f(float x) { return x * x; } ``` What should its derivative program look like? We the need the output $f(x)$ and the product of derivative at $x$, $Df(x)$ with the differential $dx$. In Slang, we put both of these together into a single function, called the *forward-mode derivative* function, which takes in a pair $(x, dx)$ returns a pair $(f(x), Df(x)\cdot dx)$ Note that in auto-diff literature, this is also often referred to as the *total derivative* function. ```C DifferentialPair fwd_f(DifferentialPair dpx) { float x = dpx.getPrimal(); // Can also be accessed via property dpx.p float dx = dpx.getDifferential(); // Can also be accessed via property dpx.d return makePair(x * x, (2 * x) * dx); } ``` Note that `(2 * x)` is the multiplier corresponding to $Df(x)$. We refer to $x$ and $f(x)$ as "*primal*" values and the perturbations $dx$ and $Df(x)\cdot dx$ as "*differential*" values. The reason for this separation is that the "*differential*" output values are always linear w.r.t their "*differential*" inputs. As the name implies, `DifferentialPair` is a special pair type used by Slang to hold values and their corresponding differentials. ### Forward-mode derivatives for higher-dimensional functions In practice, most functions tend to have multiple inputs and multiple outputs, i.e. $f: \mathbb{R}^N \to \mathbb{R}^M$ The definition above can be extended to higher dimensions, using the closest-linear-approximation idea. The main difference is that the derivative function represents a hyperplane rather than a line. Effectively, we want our forward-mode derivative to compute the following: $$ f(\mathbf{x} + \mathbf{dx}) \approx f(\mathbf{x}) + \langle Df(\mathbf{x}),\mathbf{dx}\rangle $$ Here, the input and its differential can be represented as a vector quantity $\mathbf{x}, \mathbf{dx} \in \mathbb{R}^N$ and the multiplier $Df(\mathbf{x})$ (also known as the *Jacobian* matrix) is a NxM matrix, and $\left\< \cdot,\cdot \right\>$ denotes the inner product (i.e. matrix-vector multiplication) Here's an example of a Slang function taking in two inputs (N=2) and generating one output (M=1) ```C // Compute length of hypotenuse. float f(float x, float y) { return sqrt(x * x + y * y); } ``` and its forward-mode derivative: ```C // Closest linear approximation at x, y DifferentialPair fwd_f(DifferentialPair dpx, DifferentialPair dpy) { float x = dpx.p; float y = dpy.p; float dx = dpx.d; float dy = dpx.d; return DifferentialPair( sqrt(x * x + y * y), // f(x, y) (x * dx + y * dy) / sqrt(x * x, y * y)); // } ``` Important note: the forward-mode function only needs to compute the inner product $\langle Df(\mathbf{x}),\mathbf{dx} \rangle$. The Jacobian matrix itself never needs to be fully materialized. This is a key design element of automatic differentiation, one which allows it to scale to huge input/output counts. ### Building Blocks: Forward-mode derivatives compose in forward order of execution. In practice, we compute forward-mode derivatives of a complex function by decomposing them into constituent functions (or in compiler-speak: instructions) and composing the forward-mode derivative of each piece in the **same** order. This is because of each forward derivative is a 'right-side' product (or product of Jacobian matrix with a vector) Here's an example of this in action (consider a complex function $h$ composed of $f$ and $g$): $$ h(\mathbf{x}) = f(g(\mathbf{x})) $$ It's forward-mode derivative is then: $$ \langle Dh(\mathbf{x}), \mathbf{dx}\rangle = \big\langle Df(\mathbf{x}), \langle Dg(\mathbf{x}), \mathbf{dx}\rangle\big\rangle $$ which is the forward-mode derivative of the outer function $f$ evaluated on the result of the forward-mode derivative of the inner function $g$. An example of this in Slang code: ```C // Compute square. float sqr(float x) { return x * x; } // Compute length of hypotenuse. float f(float x, float y) { float x_sqr = sqr(x); float y_sqr = sqr(y) return sqrt(x_sqr + y_sqr); } ``` The resulting derivative of `f` can be computed by composition: ```C // Forward-mode derivative of sqr() DifferentialPair fwd_sqr(DifferentialPair dpx) { float x = dpx.getPrimal(); float dx = dpx.getDifferential(); return DifferentialPair(x * x, 2 * x * dx); } // Forward-mode derivative of f() DifferentialPair fwd_f(DifferentialPair dpx, DifferentialPair dpy) { DifferentialPair dp_x_sqr = fwd_sqr(dpx); DifferentialPair dp_y_sqr = fwd_sqr(dpy); float x_sqr = dp_x_sqr.getPrimal(); float y_sqr = dp_y_sqr.getPrimal(); float x_sqr_d = dp_x_sqr.getDifferential(); float y_sqr_d = dp_y_sqr.getDifferential(); return DifferentialPair( sqrt(x_sqr + y_sqr), (x_sqr_d + y_sqr_d) / sqrt(x_sqr + y_sqr)); } ``` ### Tip: Extracting partial derivatives from a forward-mode derivative (i.e. a 'total' derivative) As we discussed above, forward-mode derivatives compute $\langle Df(\mathbf{x}),\mathbf{dx}\rangle$ rather than what you may be used to seeing in a calculus course (e.g. partial derivatives like $\frac{\partial f}{\partial x}$). In fact, the forward-mode derivative is simply an product of the partial derivative w.r.t each input parameter multiplied by their differential perturbations $\frac{\partial f}{\partial x} * dx + \frac{\partial f}{\partial x} * dy$. This is the reason for the alternative name: *total derivative*. Thus, partial derivative can be obtained by successively setting each input's differential to 1 (and 0 for everything else) Example: ```C // Compute partial derivative w.r.t x (pass dx=1.0) float df_dx = fwd_f(DifferentialPair(x, 1.0), DifferentialPair(y, 0.0)).d; // Compute partial derivaive w.r.t y (pass dy=1.0) float df_dy = fwd_f(DifferentialPair(x, 0.0), DifferentialPair(y, 1.0)).d; ``` ### Tip: Testing forward-mode derivatives using the first principles of calculus (i.e. the *finite difference* method) In Calculus, partial derivatives of a function are often defined in a 'black box' manner using limits, by perturbing a single parameter by an infinitesimal amount: $$ \frac{\partial f}{\partial x} = \lim_{dx\to 0} \frac{f(x + dx) - f(x - dx)}{2 * dx} $$ At the moment, we cannot leverage programming languages to compute true inifinitesimal limits, but we can replace $dx \to 0$ with a sufficiently small $\epsilon$ leading to the following 'test' to check if derivatives produced by automatic differentiation match with their true mathematical expected values. Here's an example of using this idea to test functions (many autodiff tests were written this way) ```C // Compute partial derivative w.r.t x analytically float df_dx_ad = fwd_f(DifferentialPair(x, 1.0), DifferentialPair(y, 0.0)) // Compute partial derivative w.r.t x through the finite difference (FD) method. float eps = 1e-4 float df_dx_fd = (f(x + eps, y) - f(x - eps, y)) / (2 * eps); // If computed correctly, df_dx_ad and df_dx_fd are very close. ``` **Caveats:** Since the finite difference method only produces a biased estimate of the derivative, the result is only numerically *close* to the auto-diff-based result. Poorly behaved functions (those that rapidly change, or are discontinuous or otherwise non-differentiable) will result in a (expected) mismatch between FD and AD results. ## Reverse-mode derivative functions This section is based off of these slides: https://cseweb.ucsd.edu/~tzli/cse291/sp2024/lectures/05_reverse_mode.pdf. ### Motivation: Challenges with scaling forward-mode derivatives A big problem with forward-mode derivatives is their inability to scale to great parameter counts. Machine learning pipelines often compute derivatives of a large complex pipeline with millions or even billions of input parameters, but a single output value, i.e. the *loss* or *objective* function, frequently denoted by $\mathcal{L}$. Computing $\frac{\partial \mathcal{L}}{\partial x_i}$ for $N$ inputs $x_i$ using the one-hot vector approach will involve invoking the forward-mode derivative function $N$ times. The reason for this limitation is that forward-mode derivatives pass derivatives from the inputs through to the outputs by computing the dot-product $\left\< Df(\mathbf{x}),\mathbf{dx}\right\>$. Instead, we employ a different approach called the reverse-mode derivative, which propagates differentials *backwards* from outputs to inputs. ### Key Idea: Generate code to compute $\langle \frac{\partial \mathcal{L}}{\partial f}, Df(\mathbf{x})\rangle$ rather than $\langle Df(\mathbf{x}),\mathbf{dx}\rangle$ The fundamental building blocks of reverse-mode derivatives are the **left-side inner product**. That is, the product of a vector of derivatives of w.r.t outputs $\frac{\partial \mathcal{L}}{\partial f}$ with the Jacobian matrix $Df(\mathbf{x})$. An important thing to keep in mind is that it does not necessarily matter what the scalar quantity $\mathcal{L}$ is. The goal of this product is to propagate the derivatives of any scalar value $\mathcal{L}$ w.r.t output vector $f(\mathbf{x})$ (i.e., $\frac{\partial \mathcal{L}}{\partial f}$) into derivatives of that same scalar value $\mathcal{L}$ w.r.t the input vector $\mathbf{x}$ (i.e., $\frac{\partial \mathcal{L}}{\partial \mathbf{x}}$). Here's an example of a Slang function computing the `reverse-mode derivative`. ```C // Compute length of hypotenuse float f(float x, float y) { return sqrt(x * x + y * y); } // Reverse-mode derivative of f. dOutput represents the derivative dL/dOutput of the output w.r.t scalar value. void rev_f(inout DifferentialPair dpx, inout DifferentialPair dpy, float dOutput) { float x = dpx.getPrimal(); float y = dpy.getPrimal(); float t = 1.0 / (sqrt(x * x + y * y)); dpx = DifferentialPair( x, // The primal part of the return value is *always* copied in from the input as-is. dOutput * x * t); // The differential part for x is the derivative dL/dx computed as // (dL/dOutput) * (dOutput/dx), where dOutput/dx = x / sqrt(x*x+y*y). dpy = DifferentialPair( y, dOutput * y * t); // The differential part for y is the derivative dL/dy computed as // (dL/dOutput) * (dOutput/dy), where dOutput/dy = y / sqrt(x*x+y*y). } ``` Note that `rev_f` accepts derivatives w.r.t the output value as the input, and returns derivatives w.r.t inputs as its output (through `inout` parameters). `rev_f` still needs the primal values `x` and `y` to compute the derivatives, so those are still passed in as an input through the primal part of the differential pair. Also note that the reverse-mode derivative function does not have to compute the primal result value (its return is void). The reason for this is a matter of convenience: reverse-mode derivatives are often invoked after all the primal functions, and there is typically no need for these values. We go into more detail on this topic in the checkpointing chapter. The reverse mode function can be used to compute both `dOutput/dx` and `dOutput/dy` with a single invocation (unlike the forward-mode case where we had to invoke `fwd_f` once for each input) ```C DifferentialPair dpx = makePair(x, 0.f); // Initialize diff-value to 0 (not necessary) DifferentialPair dpx = makePair(y, 0.f); // Initialize diff-value to 0 (not necessary) rev_f(dpx, dpy, 1.0); // Pass 1.0 for dL/dOutput so that the results are (1.0 * dOutput/dx) and (1.0 * dOutput/dy) float doutput_dx = dpx.getDifferential(); float doutput_dy = dpy.getDifferential(); ``` ### Extension to multiple outputs The extension to multiple outputs is fairly natural. Each output gets a separate input for its derivative. Here is an example: ```C // Computation involving multiple inputs and outputs. float2 f_multi_output(float x, float y) { return float2( x * x, x + y); } // Reverse-mode derivative of 'f_multi_output'. The derivative of the outputs is also a vector quantity // (type follows from return type of f_multi_output) void rev_f_multi_output(DifferentialPair dpx, DifferentialPair dpy, float2 dOut) { float x = dpx.getPrimal(); float y = dpy.getPrimal(); dpx = DifferentialPair(x, dOut[0] * 2 * x + dOut[1]); dpy = DifferentialPair(x, dOut[1]); } ``` ### Jacobian method: Generate forward- and reverse-mode derivatives from first principles. A simple way to figure out what the generated reverse (or forward) derivative function is supposed to compute is to write down the entire Jacobian function. That is, write down the partial derivative of each input w.r.t each output $$ D\mathbf{f}(\mathbf{x}) = \begin{bmatrix} \partial f_0 / \partial x & \partial f_0 / \partial y \\ \partial f_1 / \partial x & \partial f_1 / \partial y \\ \end{bmatrix} = \begin{bmatrix} 2x & 0.0 \\ 1.0 & 1.0 \\ \end{bmatrix} $$ The **reverse-mode derivative**'s outputs should match the left-product of this matrix with the vector of derivatives w.r.t outputs: $$ \left\langle \frac{\partial \mathcal{L}}{\partial \mathbf{f}}, D\mathbf{f}(\mathbf{x})\right\rangle = \begin{bmatrix} \frac{\partial \mathcal{L}}{\partial f_0} & \frac{\partial \mathcal{L}}{\partial f_1} \end{bmatrix} \begin{bmatrix} 2x & 0.0 \\ 1.0 & 1.0 \\ \end{bmatrix} = \begin{bmatrix} \left(\frac{\partial \mathcal{L}}{\partial f_0} \cdot 2x + \frac{\partial \mathcal{L}}{\partial f_1}\right) & \frac{\partial \mathcal{L}}{\partial f_1} \end{bmatrix} $$ and the **forward-mode derivative**'s outputs should match the right-product of this matrix with the vector of differentials of the inputs: $$ \langle D\mathbf{f}(\mathbf{x}), d\mathbf{x}\rangle = \begin{bmatrix} 2x & 0.0 \\ 1.0 & 1.0 \\ \end{bmatrix} \begin{bmatrix} dx \\ dy \end{bmatrix} = \begin{bmatrix} 2x \cdot dx & dx + dy \end{bmatrix} $$ Note that when we generate derivative code in practice, we do not materialize the full Jacobian matrix, and instead use the composition property to chain together derivatives at the instruction level. However, the resulting code is equivalent to the Jacobian method (mathematically), and it is a good, analytical way to confirm that the generated code is indeed correct (or when thinking about what the derivative of a particular instruction/set of instructions should be) ### Building Blocks: Reverse-mode derivatives compose in reverse order of execution. A consequence of using the 'left-side inner product' is that derivatives of a composite function must be computed in the reverse of the order of primal computation. Here's an example of a composite function $h$ (similar to the example used in forward-mode building blocks): $$ h(\mathbf{x}) = f(g(\mathbf{x})) $$ where (for brevity): $$ \mathbf{y} = g(\mathbf{x}) $$ The reverse-mode derivative function for $h$ can be written as the composition of the reverse-mode derivatives of $f$ and $g$ $$ \left\langle \frac{\partial L}{\partial h}, Dh(\mathbf{x})\right\rangle = \left\langle \left\langle \frac{\partial L}{\partial h}, Df(\mathbf{y})\right\rangle , Dg(\mathbf{x})\right\rangle $$ Note the 'backward' order here. We must first pass the derivatives through the outer function $f$, and then pass the result through the inner function $g$ to compute derivatives w.r.t inner-most inputs $\mathbf{x}$. This process of passing derivatives backwards is often referred to as *backpropagation*. A more concrete Slang example of the same: ```C // Compute square float sqr(float x) { return x * x; } // Compute length of hypotenuse float f(float x, float y) { return sqrt(sqr(x) + sqr(y)); } ``` The derivative functions are then: ```C void rev_sqr(DifferentialPair dpx, float dOutput) { float x = dpx.getPrimal(); dpx = DifferentialPair(x, dOutput * 2 * x); } void rev_f(DifferentialPair dpx, DifferentialPair dpy, float dOut) { float t = 0.5f / sqrt(x * x + y * y); float d_xsqr = t * dOut; // Calculate derivatives w.r.t output of sqr(x) float d_ysqr = t * dOut; // Calculate derivatives w.r.t output of sqr(y) rev_sqr(dpx, d_xsqr); // Propagate to x rev_sqr(dpx, d_ysqr); // Propagate to y } ``` When comparing `rev_f`'s implementation to `fwd_f`, note the order of computing derivative w.r.t `sqr` (in `rev_f`, `rev_sqr` is called at the end, while in `fwd_f` it is called at the beginning) --- ## File: docs/design/autodiff/decorators.md This document details auto-diff-related decorations that are lowered in to the IR to help annotate methods with relevant information. ## `[Differentiable]` The `[Differentiable]` attribute is used to mark functions as being differentiable. The auto-diff process will only touch functions that are marked explicitly as `[Differentiable]`. All other functions are considered non-differentiable and calls to such functions from a differentiable function are simply copied as-is with no transformation. Further, only `[Differentiable]` methods are checked during the derivative data-flow pass. This decorator is translated into `BackwardDifferentiableAttribute` (which implies both forward and backward differentiability), and then lowered into the IR `OpBackwardDifferentiableDecoration` **Note:** `[Differentiable]` was previously implemented as two separate decorators `[ForwardDifferentiable]` and `[BackwardDifferentiable]` to denote differentiability with each type of auto-diff transformation. However, these are now **deprecated**. The preferred approach is to use only `[Differentiable]` `fwd_diff` and `bwd_diff` cannot be directly called on methods that don't have the `[Differentiable]` tag (will result in an error). If non-`[Differentiable]` methods are called from within a `[Differentiable]` method, they must be wrapped in `no_diff()` operation (enforced by the [derivative data-flow analysis pass](./types.md#derivative-data-flow-analysis) ) ### `[Differentiable]` for `interface` Requirements The `[Differentiable]` attribute can also be used to decorate interface requirements. In this case, the attribute is handled in a slightly different manner, since we do not have access to the concrete implementations. The process is roughly as follows: 1. During the semantic checking step, when checking a method that is an interface requirement (in `checkCallableDeclCommon` in `slang-check-decl.cpp`), we check if the method has a `[Differentiable]` attribute 2. If yes, we construct create a set of new method declarations, one for the forward-mode derivative (`ForwardDerivativeRequirementDecl`) and one for the reverse-mode derivative (`BackwardDerivativeRequirementDecl`), with the appropriate translated function types and insert them into the same interface. 3. Insert a new member into the original method to reference the new declarations (`DerivativeRequirementReferenceDecl`) 4. When lowering to IR, the `DerivativeRequirementReferenceDecl` member is converted into a custom derivative reference by adding the `OpBackwardDerivativeDecoration(deriv-fn-req-key)` and `OpForwardDerivativeDecoration(deriv-fn-req-key)` decorations on the primal method's requirement key. Here is an example of what this would look like: ```C interface IFoo { [Differentiable] float bar(float); }; // After checking & lowering interface IFoo_after_checking_and_lowering { [BackwardDerivative(bar_bwd)] [ForwardDerivative(bar_fwd)] float bar(float); void bar_bwd(inout DifferentialPair, float); DifferentialPair bar_fwd(DifferentialPair); }; ``` **Note:** All conforming types must _also_ declare their corresponding implementations as differentiable so that their derivative implementations are synthesized to match the interface signature. In this sense, the `[Differentiable]` attribute is part of the functions signature, so a `[Differentiable]` interface requirement can only be satisfied by a `[Differentiable]` function implementation ### `[TreatAsDifferentiable]` In large codebases where some interfaces may have several possible implementations, it may not be reasonable to have to mark all possible implementations with `[Differentiable]`, especially if certain implementations use hacks or workarounds that need additional consideration before they can be marked `[Differentiable]` In such cases, we provide the `[TreatAsDifferentiable]` decoration (AST node: `TreatAsDifferentiableAttribute`, IR: `OpTreatAsDifferentiableDecoration`), which instructs the auto-diff passes to construct an 'empty' function that returns a 0 (or 0-equivalent) for the derivative values. This allows the signature of a `[TreatAsDifferentiable]` function to match a `[Differentiable]` requirement without actually having to produce a derivative. ## Custom derivative decorators In many cases, it is desirable to manually specify the derivative code for a method rather than let the auto-diff pass synthesize it from the method body. This is usually desirable if: 1. The body of the method is too complex, and there is a simpler, mathematically equivalent way to compute the same value (often the case for intrinsics like `sin(x)`, `arccos(x)`, etc..) 2. The method involves global/shared memory accesses, and synthesized derivative code may cause race conditions or be very slow due to overuse of synchronization. For this reason Slang assumes global memory accesses are non-differentiable by default, and requires that the user (or the core module) define separate accessors with different derivative semantics. The Slang front-end provides two sets of decorators to facilitate this: 1. To reference a custom derivative function from a primal function: `[ForwardDerivative(fn)]` and `[BackwardDerivative(fn)]` (AST Nodes: `ForwardDerivativeAttribute`/`BackwardDerivativeAttribute`, IR: `OpForwardDervativeDecoration`/`OpBackwardDerivativeDecoration`), and 2. To reference a primal function from its custom derivative function: `[ForwardDerivativeOf(fn)]` and `[BackwardDerivativeOf(fn)]` (AST Nodes: `ForwardDerivativeAttributeOf`/`BackwardDerivativeAttributeOf`). These attributes are useful to provide custom derivatives for existing methods in a different file without having to edit/change that module. For instance, we use `diff.meta.slang` to provide derivatives for the core module functions in `hlsl.meta.slang`. When lowering to IR, these references are placed on the target (primal function). That way both sets of decorations are lowered on the primal function. These decorators also work on generically defined methods, as well as struct methods. Similar to how function calls work, these decorators also work on overloaded methods (and reuse the `ResolveInoke` infrastructure to perform resolution) ### Checking custom derivative signatures To ensure that the user-provided derivatives agree with the expected signature, as well as resolve the appropriate method when multiple overloads are available, we check the signature of the custom derivative function against the translated version of the primal function. This currently occurs in `checkDerivativeAttribute()`/`checkDerivativeOfAttribute()`. The checking process re-uses existing infrastructure from `ResolveInvoke`, by constructing a temporary invoke expr to call the user-provided derivative using a set of 'imaginary' arguments according to the translated type of the primal method. If `ResolveInvoke` is successful, the provided derivative signature is considered to be a match. This approach also automatically allows us to resolve overloaded methods, account for generic types and type coercion. ## `[PrimalSubstitute(fn)]` and `[PrimalSubstituteOf(fn)]` In some cases, we face the opposite problem that inspired custom derivatives. That is, we want the compiler to auto-synthesize the derivative from the function body, but there _is_ no function body to translate. This frequently occurs with hardware intrinsic operations that are lowered into special op-codes that map to hardware units, such as texture sampling & interpolation operations. However, these operations do have reference 'software' implementations which can be used to produce the derivative. To allow user code to use the fast hardware intrinsics for the primal pass, but use synthesized derivatives for the derivative pass, we provide decorators `[PrimalSubstitute(ref-fn)]` and `[PrimalSubstituteOf(orig-fn)]` (AST Node: `PrimalSubstituteAttribute`/`PrimalSubstituteOfAttribute`, IR: `OpPrimalSubstituteDecoration`), that can be used to provide a reference implementation for the auto-diff pass. Example: ```C [PrimalSubstitute(sampleTexture_ref)] float sampleTexture(TexHandle2D tex, float2 uv) { // Hardware intrinsics } float sampleTexture_ref(TexHandle2D tex, float2 uv) { // Reference SW implementation. } void sampleTexture_bwd(TexHandle2D tex, inout DifferentialPair dp_uv, float dOut) { // Backward derivate code synthesized using the reference implementation. } ``` The implementation of `[PrimalSubstitute(fn)]` is relatively straightforward. When the transcribers are asked to synthesize a derivative of a function, they check for a `OpPrimalSubstituteDecoration`, and swap the current function out for the substitute function before proceeding with derivative synthesis. --- ## File: docs/design/autodiff/ir-overview.md This documentation is intended for Slang contributors and is written from a compiler engineering point of view. For Slang users, see the user-guide at this link: [https://shader-slang.com/slang/user-guide/autodiff.html](https://shader-slang.com/slang/user-guide/autodiff.html) # Overview of Automatic Differentiation's IR Passes In this document we will detail how Slang's auto-diff passes generate valid forward-mode and reverse-mode derivative functions. Refer to [Basics](./basics.md) for a review of the two derivative propagation methods and their mathematical connotations & [Types](./types.md) for a review of how types are handled under differentiation. ## Auto-Diff Pass Invocation Note that without an explicit auto-diff instruction (`fwd_diff(fn)` or `bwd_diff(fn)`) from the user present anywhere in the code, none of the auto-diff passes will do anything. Auto-diff processing operates on a function-by-function basis. Most of the logic is contained in `AutoDiffPass::processReferencedFunctions`. Here is a high-level workflow: 1. Scanning reachable insts in the module looking for `IRForwardDifferentiate` or `IRBackwardDifferentiate` operations. These instructions are added onto a work-list. The subject of a differentiate inst may be a plain function (`IRFunc`), a specialize inst (`IRSpecialize(a : IRGeneric, ...)`) in case of a generic method, or a lookup inst (`IRLookupWitness(a : IRWitnessTableType)`) when differentiating a method of an interface. 2. Dispatch each differentiation request through the appropriate 'transcriber' class. A transcriber (implements `AutodiffTranscriberBase`) is responsible for accepting a differentiation request and resolving it by replacing it with a generated function or a call to an already existing function that computes its derivative. 3. Once all currently available derivative insts have been dispatched, the follow-up work-list is checked for more transcription requests. This is a global list that all transcribers can add more follow-up work to. As an example, differentiating a function that calls another function will generate a follow-up task for this inner function, even though the latter never appears directly in a `IRForwardDifferentiate` or `IRBackwardDifferentiate` inst. At this step, there are 2 other variants that can appear `IRBackwardDifferentiatePrimal` and `IRBackwardDifferentiatePropagate` (though these can't be invoked by the user directly). 4. This process from (1.) is run in a loop. This is because we can have nested differentiation requests such as `IRForwardDifferentiate(IRBackwardDifferentiate(a : IRFuncType))`. The inner request is processed in the first pass, and the outer request gets processed in the next pass. ## Auto-Diff Passes for `IRForwardDifferentiate` For forward-mode derivatives, we only require a single pass implemented wholly in `ForwardDiffTranscriber`. This implements the linearization algorithm, which roughly follows this logic: 1. Create a clone of the original function 2. Perform pre-autodiff transformations, the most a. **Temp-Var-For-Mutable-Params** Using `IRVar` to load from b. **Linkage-Removal**: This is simply so the cloned function can be eliminated by DCE after auto-diff is complete c. **Force-Inline**: Inline all `__unsafeForceEarlyInline` & `[ForceInline]` functions _prior_ to auto-diff, so their contents can be included in the differentiation pass (even if they aren't actually marked as `[Differentiable]`) 3. Create a new blank function for the fwd-mode func (usually named `s_fwd_`) with the function type derived by transforming the original function type (See [Types](./types.md) for more information). 4. Create new blocks into the new fwd-func for each block in the orig-func. 5. Go through instructions in each block and dispatch to the appropriate generator function to emit the derivative logic into the corresponding block in the fwd-func. Each generator method is responsible for cloning in the original instruction as well as emitting a corresponding derivative instruction. `mapPrimalInst` and `mapDifferentialInst` are used to keep track of prior results so that operands for new instructions can be looked up. The generator for each instruction computes the forward-mode derivative of each *instruction* using the Jacobian method that is detailed in [Basics#Jacobian-Method](./basics.md#jacobian-method-generate-forward--and-reverse-mode-derivatives-from-first-principles). Since forward-mode derivatives can be composed in the same order as the original instructions, our generation process goes through instructions in each block in the order that they appear, creating differential insts which act as operands for future insts. Here's an example of this in IR-form ``` /* Detailed source-code truncated for AI context efficiency. */ ``` **Multiple Differential Insts:** In the above example, the derivative of each inst was a single inst. This is not always the case. For instance, `OpMul %a %b` translates to **three** insts: ```Rust %1 = OpMul %a_diff %b_primal : %float %2 = OpMul %a_primal %b_diff : %float %3 = OpAdd %1 %2 : %float ``` **Combined Primal & Differential Insts:** In some cases, there is not need to clone in the primal inst since both the primal and differential can be computed in a single inst. An example is `IRCall`, whose deriavative only needs a single call (though it needs plenty of insts to pair and unpair arguments) ```Rust // Original inst %1 = OpCall %func %a %b : %float // // Upon differentiation: // Pack args into pairs %a_pair = OpDifferentialPairMakePair %a_primal %a_diff : %dpfloat %b_pair = OpDifferentialPairMakePair %b_primal %b_diff : %dpfloat // Call into fwd-mode deriv which computes *both* primal and differential // values. // %func_fwd = OpForwardDifferentiate %func : %functype_fwd %1_pair = OpCall %func_fwd %a_pair %b_pair : %float // Split into primal and differential so they can be used for future insts. %1_primal = OpDifferentialPairGetPrimal %1_pair : %float %1_diff = OpDifferentialPairGetDifferential %1_pair : %float ``` ### Phi Arguments Block arguments are handled the same way as function arguments (which in the Slang IR, are also simply block arguments of the first block), and are converted into pair type arguments, with `OpDifferentialPairGetPrimal` and `OpDifferentialPairGetDifferential` insts automatically added to extract the primal and differential parts of each argument. ## Auto-Diff Passes for `IRBackwardDifferentiate` For reverse-mode derivatives, we need several passes that also includes differentiating the forward-mode derivative. Most of this logic is contained in `BackwardDiffTranscriberBase::transcribeFuncImpl`. These passes are inspired by the paper ["You Only Linearize Once: Tangents Transpose to Gradients"](https://arxiv.org/abs/2204.10923), which describes this approach in a functional language setting. These passes extend these ideas to work for a general-purpose imperative language structure. ### 1. Preparation The reverse-mode derivative generation involves a lot of large scale control-flow manipulation, including a CFG reversal step that aims to construct a method that flows from the end of the function to the beginning in order to compose reverse-mode derivatives. To avoid having to deal with too many corner cases (and the maintainability issues that come with it), we bring the function to a 'normal form' before running our differentiation steps. This greatly simplifies the logic of the future passes. Another high-level goal of these transformations is to bring the control-flow graph to a **reversible** form. That is, we can represent the reverse of control-flow graph using existing Slang constructs (`IRIfElse`, `IRUnconditionalBranch`, `IRLoop` and `IRSwitch`). This is not necessarily true of any valid Slang IR, so we perform additional transformations. Note: These transformations are always applied onto a temporary clone of the original function. The original function is never touched so as to not affect its use in non-autodiff contexts. Specifically we: 1. Bring the function into **single-return form**: If there are multiple blocks with return statements (i.e. multiple exit points) in a function, we eliminate this by wrapping the complete function body in a trivial loop (i.e. a single-iteration loop) and replacing existing return statements with breaks (or multi-level breaks) into its break block, which serves as the unique exit point for the function. This pass is currently contained in `convertFuncToSingleReturnForm()` 2. Eliminate **continue** statements: Loop continue statements introduce a reversibility problem. Since the forward loop can have multiple exit point, the reverse loop needs to have multiple entry points. Slang's loops do not support this. So, we eliminate these statements wrapping the body of the loop in another trivial loop (i.e. single-iteration loop) and turning the **continue** statements into **break** statements. This also involves writing **break** statements in the original loop into **multi-level** breaks. Here is an example: ```C // Original loop for (uint i = 0; i < N; i++) { if (i > 5) continue; if (i > 9) break; x = x + i; } // After continue-elimination outer_for: for (uint i = 0; i < N; i++) { inner_for: for (;;) { if (i > 5) break; if (i > 9) break outer_for; // multi-level break x = x + i; break; } } ``` 3. Eliminate **multi-level breaks**: Slang supports breaking out to an outer loop. Unfortunately, this operation is hard to reverse since Slang (and shading languages in general) do not support arbitrary `goto` statements. We eliminate multi-level breaks by assigning each nested loop a nesting index (a constant `uint` denoting the nesting level). All break statements are rewritten to break out to the immediate next level (i.e. a standard break) with a index parameter denoting the intended break level. This parameter is checked at each level and if the break index does not match the level index, we break again to the immediate upper level. This pass is currently contained in `eliminateMultiLevelBreakForFunc` Continuing the above example, here is the code after multi-level break elimination. ```C // After multi-level-break elimination uint level = -1; for (uint i = 0; i < N; i++) { for (;;) { if (i > 5) { level = 1; break; } if (i > 9) { level = 0; break; } x = x + i; level = 1; break; } if (level != 1) // Level check immediately after breaking out of each loop. break; } ``` 4. Eliminate **break** statements (enclosed in `normalizeCFG()`): Break statements also pose the same problem as continue statements (i.e. multiple exit points require the reverse loop to have multiple entry points, and Slang does not have a primitive for this). We eliminate break statements by introducing a boolean break flag which is set to `false` to indicate a break instead of using the break statement. Each *region* is enclosed in a if-else statement that checks the break flag and skips to the end if necessary. Break elimination proceeds with the following steps; Here is the above example code after break elimination. ```C // After break elimination uint level = -1; bool bflag_0 = true; // for outer loop (true => keep-going, false => break) for (uint i = 0; (i < N) && bflag_0; i++) // Insert flag into the loop condition (&& with the current condition) { bool bflag_1 = true; // for inner loop (true => keep-going, false => break) for (;bflag_1;) // Insert flag into the loop condition { if (i > 5) { level = 1; bflag_1 = false; // break } // Region after any break statement is enclosed in a // if-else check. // if (bflag_1) { if (i > 9) { level = 0; bflag_1 = false; // break } // Another if-else enclosure, this time for the second // break. // if (bflag) { x = x + i; level = 1; } bflag_1 = false; } } if (level != 1) { bflag_0 = false; } } ``` **Extra evaluation of the condition block:** The CFG normalization passes always attempt to preserve the equivalence of the original function while manipulating the control-flow constructs (i.e. ensure that the transformed code always computes the same thing). However, there is one corner-case exception: after break-elimination, the loop condition code can be evaluated 1 additional time, since we don't directly break out of the loop, but go through an extra loop condition check. This becomes important during the checkpointing step, when arrays are allocated to hold loop variables. The array bounds must account for an additional loop iteration to avoid correctness problems. ### 2. Linearization with Inst-Tagging This is the same as generating the forward-derivative function, and is in-fact handled in the same way, by invoking `ForwardDiffTranscriber`. The **inst-tagging** part of this pass is not necessary for forward-mode auto-diff (simply discarded after the auto-diff pass), but is essential for reverse-mode. **Inst-Tagging:** This pass also **tags** every instruction and block with either `IRPrimalInstDecoration`, `IRDifferentialInstDecoration` or `IRMixedDifferentialnstDecoration`, depending on whether an instruction contains/computes/reads/writes a primal value, a differential value or both. This assignment is according to the following rules: 1. The result of `.getDifferential()` from an inst of `IRDifferentialPairType` is a *differential* inst and `.getPrimal()` is a primal inst **NOTE:** This does not apply to `IRDifferentialPairUserCodeType`, all of whose operations yield a *primal* inst. 2. Further, any inst which contains a differential inst as an operand **AND** whose output value may be affected by this operand is a differential inst (e.g. if `isDifferentialInst(a) = true` then `isDifferentialInst( IRMul(a, b) ) = true`) 3. If an inst contains multiple outputs, *some* of which are differential and the others are primal, then these are *mixed-differential* insts. E.g. (a value of `IRDifferentialPairType` contains both a primal and differential value, and similarly a call of the form `IRCall(IRForwardDifferentiate(inner_fn))(...)` results in a mixed differential type since the primal part is not affect by differential inputs) 4. All other insts are *primal* by default. 5. Blocks are marked differential or primal if they contain **ONLY** differential or primal insts (respectively). Otherwise they are marked mixed-differential. The vast majority of blocks are mixed-differential. Correct tag information is critical for the next steps to correctly transform the forward-mode derivative into the reverse-mode derivative function. Here's the same forward-mode example, but with insts tagged accordingly ```Rust OpModule { // Generated function type ... // Generated function ... [OpMixedDifferentiaInstDecoration] %b_fwd = OpBlock { // Block params are mixed differentials since they carry both // primal and differential values // [OpMixedDifferentialInstDecoration] %dpa = OpParam : %dpfloat [OpMixedDifferentialInstDecoration] %dpb = OpParam : %dpfloat [OpPrimalInstDecoration] %a = OpDifferentialPairGetPrimal %dpa : %float [OpDifferentialInstDecoration] %da = OpDifferentialPairGetDifferential %dpa : %float [OpPrimalInstDecoration] %b = OpDifferentialPairGetPrimal %dpb : %float [OpDifferentialInstDecoration] %db = OpDifferentialPairGetDifferential %dpb : %float [OpPrimalInstDecoration] %1_primal = OpAdd %a %b : %float [OpDifferentialInstDecoration] %1_diff = OpAdd %da %db : %float [OpPrimalInstDecoration] %2_primal = OpAdd %1_primal %1_primal : %float [OpDifferentialInstDecoration] %2_diff = OpAdd %1_diff %1_diff : %float // Return both the primal and differential [OpMixedDifferentialInstDecoration] %2_pair = OpDifferentialPairMakePair %2_primal %2_diff : %dpfloat [OpDifferentialInstDecoration] OpReturn %2_pair } ... } ``` ### 3. Unzipping Implemented by `DiffUnzipPass`, this pass is responsible for **separating** primal instructions from differential instructions (as denoted by their decorations), by creating a full set of duplicate blocks that start **after** the last block, i.e. return block (the return statement is removed). This separation is possible because the computation of a differential inst may include primal operands but a primal inst can never use a differential operand. The unzipping pass uses the decorations from the linearization step to figure out which instructions need to be moved. The separation process uses the following high-level logic: 1. Create two clones of all the blocks in the provided function (one for primal insts, one for differential insts), and hold a mapping between each original (mixed) block to each primal and differential block. The return statement of the current final block is **removed**. 2. Process each instruction of each block: instructions marked as **primal** are moved to the corresponding **primal block**, instructions marked **differential** are moved to the corresponding **differential block**. 3. Instructions marked **mixed** need op-specific handling, and so are dispatched to the appropriate splitting function. For instance, block parameters that are holding differential-pair values are split into parameters for holding primal and differential values (the exception is function parameters, which are not affected). Similarly, `IRVar`s, `IRTerminatorInst`s (control-flow) and `IRCall`s are all split into multiple insts. 4. Except for `IRReturn`, all other control-flow insts are effectively duplicated so that the control-flow between the primal blocks and differential blocks both follow the original blocks' control-flow. The main difference is that PHI arguments are split (primal blocks carry primal values in their PHI arguments, and differential blocks carry diff values) between the two. Note that condition values (i.e. booleans) are used by both the primal and differential control-flow insts. However, since booleans are always primal values, they are always defined in the primal blocks. **Block-Tagging:** Blocks are now tagged primal or differential depending on whether they are holding primal or differential insts. This is important for the next step (transposition) to figure out which blocks need to be transposed. **Out-of-Scope Accesses:** After unzipping, the resulting IR is often **not valid**. If the control-flow is straight line (i.e. no branching or loops), the resulting IR is valid. However, if there is control-flow, then instructions can use operands whose definition does not dominate the use. This invalid IR is currently allowed to persist until the end of the auto-diff passes, when the checkpointing step occurs (i.e. Running IR validation will fail in between these steps) Here is an example of unzipped code: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` ### 4. Transposition The next step involves converting each differential instruction into its transpose. Effectively, we are re-writing each forward-mode derivative into its reverse-mode equivalent. Recall from auto-diff [basics](./basics.md), that both the forward and reverse mode derivatives can be derived from the Jacobian matrix of any operation. The main difference is whether we multiply the derivatives of the inputs with the Jacobian or multiply the Jacobian with the derivatives w.r.t the outputs. These two operations are the transpose of each other, in that the reverse-mode derivative can be thought of as multiplying with the transpose of the Jacobian. We perform this transposition on a per-instruction level. Here is an example of a transposition of a multiplication operation: ```Rust [OpPrimalInstDecoration] %b = OpLoad %var_b // %b is a primal value [OpDifferentialInstDecoration] %da = OpLoad %var_da // %da is a differential value // The operation we want to transpose [OpDifferentialInstDecoration] %1d = OpMul %da %b : %float [OpDifferentialInstDecoration] OpStore %1d %var_result ``` This multiplication can be represented as a tiny matrix multiplication between a singleton vector `[%da]` and singleton matrix `[%b]`. It's transpose will be the multiplication of the transpose of that matrix (which is the value itself `[%b]`) with a derivative w.r.t its output `%1d`, i.e. it becomes `%da = OpMul %1d %b`. Note that we now have to provide `%1d` as an **input**, and receive `da` was an output. The resulting code is then: ```Rust [OpPrimalInstDecoration] %b = OpLoad %var_b : %float // primal values are unaffected (at this stage, they are in primal blocks) // Reverse-mode code: (_rev) appended to all variables & insts to keep them distinct from the fwd-mode code. [OpDifferentialInstDecoration] %1d_rev = OpLoad %var_result_rev : %float // The operation we want to transpose [OpDifferentialInstDecoration] %da_rev = OpMul %1d_rev %b : %float [OpDifferentialInstDecoration] OpStore %da_rev %var_da_rev ``` Notice that the three differential instructions are effectively run backwards **and** transposed. Loads become stores, the `OpMul` is transposed into another `OpMul`, and stores become loads. This backwards transposition is because the differential outputs become differential inputs, and thus, we need to process the future instructions first so that the new operands are defined before bring used for the new instruction. This reverse order of operations also applies to control-flow. The rule of thumb is: if the forward-mode pass takes a particular path through the code, for a given set of primal values, the reverse-mode must "re-trace" the same path through the code, but in reverse by starting at the end. We synthesize a CFG that satisfies this property through the following steps: 1. Clone the provided unzipped forward-mode function (and all blocks + instructions) to serve as the reverse-mode function. 2. Remove all **differential** blocks and create a set of corresponding reverse-mode blocks for each **differential** block removed (**primal** blocks are simply left alone), while holding a map between corresponding blocks. Initially, they are empty. 3. Using the provided unzipped forward-mode function as a reference, process each differential block by walking each instruction from the _last_ (terminator) inst, and dispatching to the appropriate op-specific `transposeXYZ()` method to emit the appropriate transposed instructions into the corresponding reverse-mode block. There are several concerns that must be taken care of: 1. **Multiple Derivaive Outputs:** Unlike forward-mode auto-diff, where an inst producing a single value, would only need a single derivative (corresponding to that value), reverse-mode auto-diff can produce multiple derivatives from an inst. For instance `%dc = IRAdd(%da, %db)` produces two derivatives: `%da_rev = %dc_rev` and `%db_rev = %dc_rev`. Thus, the `transposeXYZ()` implementation for any instruction can return a set of derivative insts for each relevant input differential value. 2. **Insts Used in Multiple Places (Derivative Accumulation):** If an inst is used in multiple places, and receives a reverse-mode derivative from several of those places, these results need to be **added up** to get the correct derivative. Consider this forward-mode example ```Rust [OpDifferentialInstDecoration] %db = OpAdd %da, %da : %float [OpDifferentialInstDecoration] %dc = OpAdd %db, %da : %float ``` It's reverse-mode derivative will look like this: ```Rust %db_rev = %dc_rev // %db only has one differential since it only consumed in one place. // reverse-mode differential for %da from trnaposing the first instruction [OpDifferentialInstDecoration] %da_rev_1 = OpAdd %db_rev, %db_rev : %float // reverse-mode differential for %da from transposing the second instruction [OpDifferentialInstDecoration] %da_rev_2 = %dc_rev // add them together to get the final derivative for %da [OpDifferentialInstDecoration] %da_rev = OpAdd %da_rev_1 %da_rev_2 : %float ``` Derivative accumulation is achieved through two ways: **Within** a block, we keep a list all the reverse derivative insts for each inst and only **materialize** the total derivative when it is required as an operand. This is the most efficient way to do this, because we can apply certain optimizations for composite types (derivative of an array element, vector element, struct field, etc..). **Across** blocks, we use an accumulator variable that is inserted into a top-level block in the function, and add to this variable whenever a transposition operation generates a new inst. This can sometimes produce sub-optimal code for aggregate/large data types, but at the moment, the accumulator method is necessary because insts can receive derivatives from conditionally executed blocks. While this example uses `OpAdd` to demonstrate accumulation, in practice, we use the derivative type system (See [Types](./types.md) for more) to look up the derivative addition function (`dadd`) to add two values of an arbitrary differential type. In practice, the `OpAdd` is replaced by `OpCall %float_dadd %da_rev1 %da_rev_2`. Similarly, for accumulator variables, we must initialize them to zero for the accumulation to work correctly, and we lookup the `dzero` interface method to initialize it in a type-specific way. 3. **Deferred Materialization for Derivatives of Composite Types:** Non-primitive types, such as vectors, arrays, structs, etc. whose elements are used in several places in the forward-mode code, can result in sub-optimal reverse-mode code. Here is an example (in Slang source-style): ```C float f_fwd(DifferentialPair input) { float3 dinput = input.getDifferential(); float a = dinput.x + dinput.y; float b = a + dinput.z; return b; } // Transposed code (naively, without deferred materialization) void f_rev(inout DifferentialPair input, float d_output) { // transpose of (return b;) float db_rev = d_output; // transpose of (float b = a + dinput.z) float da_rev = db_rev; float3 dinput_rev_1 = float3(0.f, 0.f, da_rev); // transpose of (float a = dinput.x + dinput.y) float3 dinput_rev_2 = float3(0.f, da_rev, 0.f); float3 dinput_rev_3 = float3(da_rev, 0.f, 0.f); // Accumulate [dinput_rev_1, dinput_rev_2, dinput_rev_3] float3 dinput = dinput_rev_1 + dinput_rev_2 + dinput_rev_3 input = DifferentialPair( input.getPrimal(), dinput); } ``` Note that, this approach to inst-by-inst transposition can use a lot more stack space than is necessary (`dinput_rev_1`, `dinput_rev_2` and `dinput_rev_3` all only have a single non-0 entry). This is a known complexity issue with naive inst-by-inst transposition: hypothetically, an size-$N$ vector/array would end up allocating $O(N^2)$ memory even if only $N$ elements are non-0. In our Slang implementation, we circumvent this (to an extent) by deferring materialization. Rather than create each component `dinput_rev_i` as soon as we see an inst use, we hold the derivative with a special flavor value for lookups (say `Swizzle` or `GetElement`). When the total value `dinput_rev` is necessary, we process components of each flavor type at once and create a single derivative from all the components. Here is the same example, with deferred materialization: ```C // Transposed code (naively, without deferred materialization) void f_rev(inout DifferentialPair input, float d_output) { // transpose of (return b;) float db_rev = d_output; // transpose of (float b = a + dinput.z), hold {flavor=Swizzle, component=.z, derivInst=db_rev} in list. float da_rev = db_rev; // transpose of (float a = dinput.x + dinput.y), // hold {flavor=Swizzle, component=.x, derivInst=da_rev} and {flavor=Swizzle, component=.y, derivInst=da_rev} in list. // Materialize when required (for constructing return pair) float3 dinput = float3(db_rev, da_rev, da_rev); input = DifferentialPair( input.getPrimal(), dinput); } ``` Note that this only really works for accumulation *within* a single block/control-flow region. For across regions, we still have to materialize when we exit a region, so this memory problem can still manifest for control-flow heavy functions, where each region must allocate enough space for its contribution to the full derivative, even if only a small subset is non-0. ```C float a[10] = /*...*/; for (int i = 0; i < 10; i++) { a[i] = f(a[i]); } ``` ```C // Entry block %t = OpBlock { IRLoop %c %br %c 0 } // Condition %c = OpBlock { %i = OpParam : %float %a = OpParam : %Array(%float, 10) %2 = OpLesser(%i, 10) : %bool %OpIfElse(%2, %b, %br, %br) } // Loop body. %b = OpBlock { %a_i = OpGetElement(%a, %i) : %float %f_a_i = OpCall(f, %a_i) : %float %a_next = OpUpdateElement(%a, %i, %f_a_i) : %Array(%float, 10) %i_next = OpAdd(%i, 1) OpUnconditionalBranch(%c, %i_next, %a_next) } // Break block %br = OpBlock { //... } ``` After AD passes, this results in the following code: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` 4. Construct the reverse control-flow (`reveseCFGRegion()`) by going through the reference forward-mode blocks, and cloning the control-flow onto the reverse-mode blocks, but in reverse. This is achieved by running `reverseCFGRegion()` recursively on each sub-region, where a *region* is defined as a set of blocks with a single entry block and a single exit block. This definition of a region only works because we normalized the CFG into this form. The reversal logic follows these general rules: 1. **Unconditional Branch**: For an unconditional branch from `A->B` we simply have to map the reverse version of B with that of A. i.e. `rev[B] -> rev[A]` 2. **If-Else**: For an if-else of the form `A->[true = T->...->T_last->M, false = F->...->F_last->M]`, we construct `rev[M]->[true = rev[T_last]->...->rev[T_last]->rev[A], false = rev[F_last]->...->rev[F]->rev[A]]`. That is, we reverse each sub-region, and start from the merge block and end at the split block. Note that we need to identify `T_last` and `F_last` i.e. the last two blocks in the true and false regions. We make the last block in the region an additional return value of `reverseCFGRegion()`, so that when reversing the true and false sub-regions, we also get the relevant last block as an additional output. Also note that additional empty blocks may be inserted to carry derivatives of the phi arguments, but this does not alter the control-flow. 3. **Switch-case**: Proceeds in exactly the same way as `if-else` reversal, but with multiple cases instead of just 2. 4. **Loop**: After normalization, all (non-trivial) loops are of the form: `A->C->[true = T->...->T_last->C, false=B->...->M]`. We reverse this loop into `rev[M]->...rev[B]->rev[C]->[true=rev[T_last]->...->rev[T]->rev[C], false=rev[A]]`. The actual reversal logic also handles some corner cases by inserting additional blank blocks to avoid situations where regions may share the same merge block. Finally, we process the first and last blocks (entry and return blocks) by inserting a void return (reverse-mode derivative functions are always of void result type) At this stage, the reverse-mode generation is almost complete. The control-flow and the derivative logic is present, but we still have to resolve out-of-scope accesses from the new differential blocks into the primal block. ### 5. Checkpointing/Recomputation (also called 'primal-hoisting') This step legalizes the out-of-scope accesses of primal insts from within differential blocks. This is to prepare us for the next step (i.e. [extraction](#6-extraction)) that splits the function into two by moving the primal blocks into a separate primal-context-generator function, and the differential blocks into the backward-propagation function. Before we can perform this extraction, we must find any primal values being used in differential blocks and handle them in one of two ways: **Store** (put the values in a static struct) or **Recompute** (clone the necessary instructions to recompute when necessary). We first _classify_ all necessary instructions into one of the two buckets before processing each use accordingly. 1. **Classify uses into each set:** Note that rather than proceeding on an inst-by-inst basis, we classify **uses** of insts. The same inst can be used in several places, and we may decide to store one use and recompute another (in some cases, this could be the optimal result). The classification process uses a work-list approach that roughly looks like the following: 1. Add all uses of **primal** insts in an inst within a **differential** block to the work list. This is our initial set of uses that require classification. 2. Query the active policy object (which for now is hardcoded) to obtain the classification based on heuristics & user decorations (Specifically `[PreferRecompute]` and `[PreferCheckpoint]` decorations influence the classification policy) 3. For uses that should be **recomputed**, we have to now make the same decision one their **operands**, in order to make them available for the recomputation insts. Thus, their operands are added to the work list. 4. For uses that should be **stored**, there is no need to consider their operands, since the computed value will be explicitly stored and loaded later. 5. Once the worklist is empty, go over all the **uses** and their classifications, and convert them into a list of **insts** that should be stored or recomputed. Note that if an inst has uses with both classifications, then it can appear in both lists. 2. **Process 'Store' (i.e. checkpoint) insts:** Store them into a single variable (of a struct type that is synthesized as necessary), and then loaded from in the differential blocks. This allows us to simply turn this variable into an output parameter from the context function and an input parameter for the backprop function. When storing values this way, we must consider that instructions within loops can have different values each iteration. Thus, we must use an array to store each value, and this array's size must be statically known since we wish to synthesize a static struct type to hold all the stored values. Thus, we enforce the requirement of a `[MaxIters(N)]` decoration and attempt to infer a loop iteration limit if one is not provided. Here's an example of a case where we decide to checkpoint _all_ relevant uses: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` Another example with a function `g` that does contain loops: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` **Indexed Region Processing:** In order to be able to allocate the right array and use the right indices, we need information about which blocks are part of which loop (and loops can be nested, so blocks can be part of multiple loops). To do this, we run a pre-processing step that maps all blocks to all relevant loop regions, the corresponding index variables and the inferred iteration limits (maximum times a loop can run). Note that if an instruction appears in a nested block, we create a multi-dimensional array and use multiple indices. **Loop State Variables:** Certain variables cannot be classified as recompute. Major examples are loop state variables which are defined as variables that are read from and written to within the loop. In practice, they appear as phi-variables on the first loop block after SSA simplification. Their uses _must_ be classified as 'store', because recomputing them requires duplicating the primal loop within the differential loop. This is because the differential loop runs backwards so the state of a primal variable at loop index $N$ cannot be recomputed when the loop is running backwards ($N+1 \to N \to N-1$), and involves running the primal loop up to $N$ times within the current iteration of the differential loop. In terms of complexity, this turns an $O(N)$ loop into an $O(N^2)$ loop, and so we disallow this. It is possible that the resulting $O(N^2)$ loop may end up being faster in practice due to reduced memory requirements, but we currently lack the infrastructure to robustly allow such loop duplication while keeping the user informed of the potentially drastic complexity issues. 3. **Process 'Recompute' insts:** Insert a copy of the primal instruction into a corresponding 'recomputation' block that is inserted into the differential control-flow so that it dominates the use-site. **Insertion of Recompute Blocks:** In order to accommodate recomputation, we first preprocess the function, by going through each **breakable (i.e. loop) region** in the differential blocks, looking up the corresponding **primal region** and cloning all the primal blocks into the beginning of the differential region. Note that this cloning process does not actually clone the instructions within each block, only the control-flow (i.e. terminator) insts. This way, there is a 1:1 mapping between the primal blocks and the newly created **recompute blocks**, This way, if we decide to 'recompute' an instruction, we can simply clone it into the corresponding recompute block, and we have a guarantee that the definition and use-site are within the same loop scope, and that the definition comes before the use. **Legalizing Accesses from Branches:** Our per-loop-region recompute blocks ensure that the recomputed inst is always within the same region as its uses, but it can still be out-of-scope if it is defined within a branch (i.e. if-else). We therefore still run a light-weight hoisting pass that detects these uses, inserts an `IRVar` at the immediate dominator of the def and use, and inserts loads and stores accordingly. Since they occur within the same loop region, there is no need to worry about arrays/indices (unlike the 'store' case). **Marking Recompute Blocks:** These blocks are marked with `OpRecomputeBlockDecoration` to identify them as containing primal instructions, even though they are within differential regions. This helps us remove any unused blocks if none of the instructions end up being recomputed. Here is an example of recomputation demonstrated in Slang source-style (although this takes place in IR-form) ``` /* Detailed source-code truncated for AI context efficiency. */ ``` For completeness, here is another example of a function `g` which contains a loop to demonstrate how recomputation works when there are multiple loop regions. ``` /* Detailed source-code truncated for AI context efficiency. */ ``` ### 6. Extraction The final step involves _splitting_ the function immediately after the primal block to create two functions: a **primal context function** that computes the primal value normally, but also outputs a context object with relevant intermediate values, and a **backward propagation function** that computes the backward derivative and consumes this context object for the required intermediate values. The first 5 steps have set us up for this final step, so it is not particularly complex. We follow this high-level logic: 1. Create an empty function for the primal context function. The type of this function is the same as the primal function, but with an additional `out` parameter for the intermediate context, whose type is undecided at this stage. We use a temporary function-specific type called `OpBackwardDerivativeIntermediateContextType(func)` as a placeholder. 2. Move primal blocks to the primal context function. Re-create the return inst (the return value is temporarily remembered using a decoration during the rest of the AD process). Also, the first block (reserved for function parameters) is also duplicated and processed to have primal parameters in the primal function and pair parameters in the differential function. 3. Lower all `OpBackwardDerivativeIntermediateContextType` types into concrete struct types by creating a field for each 'stored' inst from Step 5. This lowering process happens **at the end of the current AD pass after all relevant methods have completed Step 5**. We need Step 5 (hoisting) to be complete for all relevant methods because the context struct for a given function can include context structs of other functions that are called from it. Our context-type lowering therefore proceeds recursively by lowering the context for inner functions as necessary. The lowering process also removes the temporary vars that were created to hold the store insts, and replaces them with a stores and loads from the context struct. **Recursive Functions are Disallowed:** Since we lower all intermediate types into a static struct type, recursive calls cannot currently be supported from differentiable functions. The context struct for a method may include itself, creating an impossible scenario. Here is one of the examples above (`g`) after checkpointing: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` Having separate methods for the primal and backward passes is necessary when reverse-mode differentiating a method that calls out to other differentiable functions. Here is an example of differentiating a method that calls out to multiple methods, to get an idea for why we need the primal context method to be separate ```C float outer(float x) { float y = f(x); float z = g(y); float w = h(z); return w; } // It's complete reverse mode derivative looks like the following: void outer_rev(DifferentialPair dpx, float d_output) { // Compute the primal values in the forward direction, while producing relevant context. f_Intermediates f_ctx; g_Intermediates g_ctx; h_Intermediates h_ctx; float y = s_primal_ctx_f(x, f_ctx); float z = s_primal_ctx_g(y, g_ctx); float w = s_primal_ctx_h(z, h_ctx); // Note that at this point, we are holding intermediate context variables for f, g and h. // Consume the context while evaluating the propagating the derivatives backwards. DifferentialPair dpz = {z, 0.f}; s_bwd_h(dpz, d_output, h_ctx); DifferentialPair dpy = {y, 0.f}; s_bwd_g(dpy, dpz.getDifferential(), g_ctx); DifferentialPair _dpx = {x, 0.f}; s_bwd_f(dpx, dpy.getDifferential(), f_ctx); dpx = _dpx; } ``` --- ## File: docs/design/autodiff/types.md This documentation is intended for Slang contributors and is written from a compiler engineering point of view. For Slang users, see the user-guide at this link: [https://shader-slang.com/slang/user-guide/autodiff.html](https://shader-slang.com/slang/user-guide/autodiff.html) Before diving into this document, please review the document on [Basics](./basics.md) for the fundamentals of automatic differentiation. # Components of the Type System Here we detail the main components of the type system: the `IDifferentiable` interface to define differentiable types, the `DifferentialPair` type to carry a primal and corresponding differential in a single type. We also detail how auto-diff operators are type-checked (the higher-order function checking system), how the `no_diff` decoration can be used to avoid differentiation through attributed types, and the derivative data flow analysis that warns the the user of unintentionally stopping derivatives. ## `interface IDifferentiable` Defined in core.meta.slang, `IDifferentiable` forms the basis for denoting differentiable types, both within the core module, and otherwise. The definition of `IDifferentiable` is designed to encapsulate the following 4 items: 1. `Differential`: The type of the differential value of the conforming type. This allows custom data-structures to be defined to carry the differential values, which may be optimized for space instead of relying solely on compiler synthesis/ Since the computation of derivatives is inherently linear, we only need access to a few operations. These are: 2. `dadd(Differential, Differential) -> Differential`: Addition of two values of the differential type. It's implementation must be associative and commutative, or the resulting derivative code may be incorrect. 3. `dzero() -> Differential`: Additive identity (i.e. the zero or empty value) that can be used to initialize variables during gradient aggregation 4. `dmul(S, Differential)`: Scalar multiplication of a real number with the differential type. It's implementation must be distributive over differential addition (`dadd`). Points 2, 3 & 4 are derived from the concept of vector spaces. The derivative values of any Slang function always form a vector space (https://en.wikipedia.org/wiki/Vector_space). ### Derivative member associations In certain scenarios, the compiler needs information on how the fields in the original type map to the differential type. Particularly, this is a problem when differentiate the implicit construction of a struct through braces (i.e. `{}`), represented by `kIROp_MakeStruct`. We provide the decorator `[DerivativeMember(DifferentialTypeName.fieldName)]` (ASTNode: DerivativeMemberAttribute, IR: kIROp_DerivativeMemberDecoration) to explicitly mark these associations. Example ```C struct MyType : IDifferentiable { typealias Differential = MyDiffType; float a; [DerivativeMember(MyDiffType.db)] float b; /* ... */ }; struct MyDiffType { float db; }; ``` ### Automatic Synthesis of `IDifferentible` Conformances for Aggregate Types It can be tedious to expect users to hand-write the associated `Differential` type, the corresponding mappings and interface methods for every user-defined `struct` type. For aggregate types, these are trivial to construct by analysing which of their components conform to `IDifferentiable`. The synthesis proceeds in roughly the following fashion: 1. `IDifferentiable`'s components are tagged with special decorator `__builtin_requirement(unique_integer_id)` which carries an enum value from `BuiltinRequirementKind`. 2. When checking that types conform to their interfaces, if a user-provided definition does not satisfy a requirement with a built-in tag, we perform synthesis by dispatching to `trySynthesizeRequirementWitness`. 3. For _user-defined types_, Differential **types** are synthesized during conformance-checking through `trySynthesizeDifferentialAssociatedTypeRequirementWitness` by checking if each constituent type conforms to `IDifferentiable`, looking up the corresponding `Differential` type, and constructing a new aggregate type from these differential types. Note that since it is possible that a `Differential` type of a constituent member has not yet been synthesized, we have additional logic in the lookup system (`trySynthesizeRequirementWitness`) that synthesizes a temporary empty type with a `ToBeSynthesizedModifier`, so that the fields can be filled in later, when the member type undergoes conformance checking. 4. For _user-defined types_, Differential methods (`dadd`, `dzero` and `dmul`) are synthesized in `trySynthesizeDifferentialMethodRequirementWitness` by utilizing the `Differential` member and its `[DifferentialMember]` decorations to determine which fields need to be considered and the base type to use for each field. There are two synthesis patterns. The fully-inductive pattern is used for `dadd` and `dzero` which works by calling `dadd` and `dzero` respectively on the individual fields of the `Differential` type under consideration. Example: ```C // Synthesized from "struct T {FT1 field1; FT2 field2;}" T.Differential dadd(T.Differential a, T.Differential b) { return Differential( FT1.dadd(a.field1, b.field1), FT2.dadd(a.field2, b.field2), ) } ``` On the other hand, `dmul` uses the fixed-first arg pattern since the first argument is a common scalar, and proceeds inductively on all the other args. Example: ```C // Synthesized from "struct T {FT1 field1; FT2 field2;}" T.Differential dmul(S s, T.Differential a) { return Differential( FT1.dmul(s, a.field1), FT2.dmul(s, a.field2), ) } ``` 5. During auto-diff, the compiler can sometimes synthesize new aggregate types. The most common case is the intermediate context type (`kIROp_BackwardDerivativeIntermediateContextType`), which is lowered into a standard struct once the auto-diff pass is complete. It is important to synthesize the `IDifferentiable` conformance for such types since they may be further differentiated (through higher-order differentiation). This implementation is contained in `fillDifferentialTypeImplementationForStruct(...)` and is roughly analogous to the AST-side synthesis. ### Differentiable Type Dictionaries During auto-diff, the IR passes frequently need to perform lookups to check if an `IRType` is differentiable, and retrieve references to the corresponding `IDifferentiable` methods. These lookups also need to work on generic parameters (that are defined inside generic containers), and existential types that are interface-typed parameters. To accommodate this range of different type systems, Slang uses a type dictionary system that associates a dictionary of relevant types with each function. This works in the following way: 1. When `CheckTerm()` is called on an expression within a function that is marked differentiable (`[Differentiable]`), we check if the resolved type conforms to `IDifferentiable`. If so, we add this type to the dictionary along with the witness to its differentiability. The dictionary is currently located on `DifferentiableAttribute` that corresponds to the `[Differentiable]` modifier. 2. When lowering to IR, we create a `DifferentiableTypeDictionaryDecoration` which holds the IR versions of all the types in the dictionary as well as a reference to their `IDifferentiable` witness tables. 3. When synthesizing the derivative code, all the transcriber passes use `DifferentiableTypeConformanceContext::setFunc()` to load the type dictionary. `DifferentiableTypeConformanceContext` then provides convenience functions to lookup differentiable types, appropriate `IDifferentiable` methods, and construct appropriate `DifferentialPair`s. ### Looking up Differential Info on _Generic_ types Generically defined types are also lowered into the differentiable type dictionary, but rather than having a concrete witness table, the witness table is itself a parameter. When auto-diff passes need to find the differential type or place a call to the IDifferentiable methods, this is turned into a lookup on the witness table parameter (i.e. `Lookup(, )`). Note that these lookups instructions are inserted into the generic parent container rather than the inner most function. Example: ```C T myFunc(T a) { return a * a; } // Reverse-mode differentiated version void bwd_myFunc( inout DifferentialPair dpa, T.Differential dOut) // T.Differential is Lookup('Differential', T_Witness_Table) { T.Differential da = T.dzero(); // T.dzero is Lookup('dzero', T_Witness_Table) da = T.dadd(dpa.p * dOut, da); // T.dadd is Lookup('dadd', T_Witness_Table) da = T.dadd(dpa.p * dOut, da); dpa = diffPair(dpa.p, da); } ``` ### Looking up Differential Info on _Existential_ types Existential types are interface-typed values, where there are multiple possible implementations at run-time. The existential type carries information about the concrete type at run-time and is effectively a 'tagged union' of all possible types. #### Differential type of an Existential The differential type of an existential type is tricky to define since our type system's only restriction on the `.Differential` type is that it also conforms to `IDifferentiable`. The differential type of any interface `IInterface : IDifferentiable` is therefore the interface type `IDifferentiable`. This is problematic since Slang generally requires a static `anyValueSize` that must be a strict upper bound on the sizes of all conforming types (since this size is used to allocate space for the union). Since `IDifferentiable` is defined in the core module `core.meta.slang` and can be used by the user, it is impossible to define a reliable bound. We instead provide a new **any-value-size inference** pass (`slang-ir-any-value-inference.h`/`slang-ir-any-value-inference.cpp`) that assembles a list of types that conform to each interface in the final linked IR and determines a relevant upper bound. This allows us to ignore types that conform to `IDifferentiable` but aren't used in the final IR, and generate a tighter upper bound. **Future work:** This approach, while functional, creates a locality problem since the size of `IDifferentiable` is the max of _all_ types that conform to `IDifferentiable` in visible modules, even though we only care about the subset of types that appear as `T.Differential` for `T : IInterface`. The reason for this problem is that upon performing an associated type lookup, the Slang IR drops all information about the base interface that the lookup starts from and only considers the constraint interface (in this case `Differential : IDifferentiable`). There are several ways to resolve this issue, including (i) a static analysis pass that determines the possible set of types at each use location and propagates them to determine a narrower set of types, or (ii) generic (or 'parameterized') interfaces, such as `IDifferentiable` where each version can have a different set of conforming types. Example: ```C interface IInterface : IDifferentiable { [Differentiable] This foo(float val); [Differentiable] float bar(); }; float myFunc(IInterface obj, float a) { IInterface k = obj.foo(a); return k.bar(); } // Reverse-mode differentiated version (in pseudo-code corresponding to IR, some of these will get lowered further) void bwd_myFunc( inout DifferentialPair dpobj, inout DifferentialPair dpa, float.Differential dOut) // T.Differential is Lookup('Differential', T_Witness_Table) { // Primal pass.. IInterface obj = dpobj.p; IInterface k = obj.foo(a); // ..... // Backward pass DifferentialPair dpk = diffPair(k); bwd_bar(dpk, dOut); IDifferentiable dk = dpk.d; // Differential of `IInterface` is `IDifferentiable` DifferentialPair dp = diffPair(dpobj.p); bwd_foo(dpobj, dpa, dk); } ``` #### Looking up `dadd()` and `dzero()` on Existential Types There are two distinct cases for lookup on an existential type. The more common case is the closed-box existential type represented simply by an interface. Every value of this type contains a type identifier & a witness table identifier along with the value itself. The less common case is when the function calls are performed directly on the value after being cast to the concrete type. **`dzero()` for "closed" Existential type: The `NullDifferential` Type** For concrete and even generic types, we can initialize a derivative accumulator variable by calling the appropriate `Type.dzero()` method. This is unfortunately not possible when initializing an existential differential (which is currently of type `IDifferentiable`), since we must also initialize the type-id of this existential to one of the implementations, but we do not know which one yet since that is a run-time value that only becomes known after the first differential value is generated. To get around this issue, we declare a special type called `NullDifferential` that acts as a "none type" for any `IDifferentiable` existential object. **`dadd()` for "closed" Existential types: `__existential_dadd`** We cannot directly use `dadd()` on two existential differentials of type `IDifferentiable` because we must handle the case where one of them is of type `NullDifferential` and `dadd()` is only defined for differentials of the same type. We handle this currently by synthesizing a special method called `__existential_dadd` (`getOrCreateExistentialDAddMethod` in `slang-ir-autodiff.cpp`) that performs a run-time type-id check to see if one of the operand is of type `NullDifferential` and returns the other operand if so. If both are non-null, we dispatch to the appropriate `dadd` for the concrete type. **`dadd()` and `dzero()` for "open" Existential types** If we are dealing with values of the concrete type (i.e. the opened value obtained through `ExtractExistentialValue(ExistentialParam)`). Then we can perform lookups in the same way we do for generic type. All existential parameters come with a witness table. We insert instructions to extract this witness table and perform lookups accordingly. That is, for `dadd()`, we use `Lookup('dadd', ExtractExistentialWitnessTable(ExistentialParam))` and place a call to the result. ## `struct DifferentialPair` The second major component is `DifferentialPair` that represents a pair of a primal value and its corresponding differential value. The differential pair is primarily used for passing & receiving derivatives from the synthesized derivative methods, as well as for block parameters on the IR-side. Both `fwd_diff(fn)` and `bwd_diff(fn)` act as function-to-function transformations, and so the Slang front-end translates the type of `fn` to its derivative version so the arguments can be type checked. ### Pair type lowering. The differential pair type is a special type throughout the AST and IR passes (AST Node: `DifferentialPairType`, IR: `kIROp_DifferentialPairType`) because of its use in front-end semantic checking and when synthesizing the derivative code for the functions. Once the auto-diff passes are complete, the pair types are lowering into simple `struct`s so they can be easily emitted (`DiffPairLoweringPass` in `slang-ir-autodiff-pairs.cpp`). We also define additional instructions for pair construction (`kIROp_MakeDifferentialPair`) and extraction (`kIROp_DifferentialPairGetDifferential` & `kIROp_DifferentialPairGetPrimal`) which are lowered into struct construction and field accessors, respectively. ### "User-code" Differential Pairs Just as we use special IR codes for differential pairs because they have special handling in the IR passes, sometimes differential pairs should be _treated as_ regular struct types during the auto-diff passes. This happens primarily during higher-order differentiation when the user wishes to differentiate the same code multiple times. Slang's auto-diff approaches this by rewriting all the relevant differential pairs into 'irrelevant' differential pairs (`kIROp_DifferentialPairUserCode`) and 'irrelevant' accessors (`kIROp_DifferentialPairGetDifferentialUserCode`, `kIROp_DifferentialPairGetPrimalUserCode`) at the end of **each auto-diff iteration** so that the next iteration treats these as regular differentiable types. The user-code versions are also lowered into `struct`s in the same way. ## Type Checking of Auto-Diff Calls (and other _higher-order_ functions) Since `fwd_diff` and `bwd_diff` are represented as higher order functions that take a function as an input and return the derivative function, the front-end semantic checking needs some notion of higher-order functions to be able to check and lower the calls into appropriate IR. ### Higher-order Invocation Base: `HigherOrderInvokeExpr` All higher order transformations derive from `HigherOrderInvokeExpr`. For auto-diff there are two possible expression classes `ForwardDifferentiateExpr` and `BackwardDifferentiateExpr`, both of which derive from this parent expression. ### Higher-order Function Call Checking: `HigherOrderInvokeExprCheckingActions` Resolving the concrete method is not a trivial issue in Slang, given its support for overloading, type coercion and more. This becomes more complex with the presence of a function transformation in the chain. For example, if we have `fwd_diff(f)(DiffPair(...), DiffPair(...))`, we would need to find the correct match for `f` based on its post-transform argument types. To facilitate this we use the following workflow: 1. The `HigherOrderInvokeExprCheckingActions` base class provides a mechanism for different higher-order expressions to implement their type translation (i.e. what is the type of the transformed function). 2. The checking mechanism passes all detected overloads for `f` through the type translation and assembles a new group out of the results (the new functions are 'temporary') 3. This new group is used by `ResolveInvoke` when performing overload resolution and type coercion using the user-provided argument list. 4. The resolved signature (if there is one) is then replaced with the corresponding function reference and wrapped in the appropriate higher-order invoke. **Example:** Let's say we have two functions with the same name `f`: (`int -> float`, `double, double -> float`) and we want to resolve `fwd_diff(f)(DiffPair(1.0, 0.0), DiffPair(0.0, 1.0))`. The higher-order checking actions will synthesize the 'temporary' group of translated signatures (`int -> DiffPair`, `DiffPair, DiffPair -> DiffPair`). Invoke resolution will then narrow this down to a single match (`DiffPair, DiffPair -> DiffPair`) by automatically casting the `float`s to `double`s. Once the resolution is complete, we return `InvokeExpr(ForwardDifferentiateExpr(f : double, double -> float), casted_args)` by wrapping the corresponding function in the corresponding higher-order expr ## Attributed Types (`no_diff` parameters) Often, it will be necessary to prevent gradients from propagating through certain parameters, for correctness reasons. For example, values representing random samples are often not differentiated since the result may be mathematically incorrect. Slang provides the `no_diff` operator to mark parameters as non-differentiable, even if they use a type that conforms to `IDifferentiable` ```C float myFunc(float a, no_diff float b) { return a * b; } // Resulting fwd-mode derivative: DiffPair myFunc(DiffPair dpa, float b) { return diffPair(dpa.p * b, dpa.d * b); } ``` Slang uses _OpAttributedType_ to denote the IR type of such parameters. For example, the lowered type of `b` in the above example is `OpAttributedType(OpFloat, OpNoDiffAttr)`. In the front-end, this is represented through the `ModifiedType` AST node. Sometimes, this additional layer can get in the way of things like type equality checks and other mechanisms where the `no_diff` is irrelevant. Thus, we provide the `unwrapAttributedType` helper to remove attributed type layers for such cases. ## Derivative Data-Flow Analysis Slang has a derivative data-flow analysis pass that is performed on a per-function basis immediately after lowering to IR and before the linking step (`slang-ir-check-differentiability.h`/`slang-ir-check-differentiability.cpp`). The job of this pass is to enforce that instructions that are of a differentiable type will propagate a derivatives, unless explicitly dropped by the user through `detach()` or `no_diff`. The reason for this is that Slang requires functions to be decorated with `[Differentiable]` to allow it to propagate derivatives. Otherwise, the function is considered non-differentiable, and effectively produces a 0 derivative. This can lead to frustrating situations where a function may be dropping non-differentiable on purpose. Example: ```C float nonDiffFunc(float x) { /* ... */ } float differentiableFunc(float x) // Forgot to annotate with [Differentiable] { /* ... */ } float main(float x) { // User doesn't realise that the function that is supposed to be differentiable is not // getting differentiated, because the types here are all 'float'. // return nonDiffFunc(x) * differentiableFunc(x); } ``` The data-flow analysis step enforces that non-differentiable functions used in a differentiable context should get their derivative dropped explicitly. That way, it is clear to the user whether a call is getting differentiated or dropped. Same example with `no_diff` enforcement: ```C float nonDiffFunc(float x) { /* ... */ } [Differentiable] float differentiableFunc(float x) { /* ... */ } float main(float x) { return no_diff(nonDiffFunc(x)) * differentiableFunc(x); } ``` A `no_diff` can only be used directly on a function call, and turns into a `TreatAsDifferentiableDecoration` that indicates that the function will not produce a derivative. The derivative data-flow analysis pass works similar to a standard data-flow pass: 1. We start by assembling a set of instructions that 'produce' derivatives by starting with the parameters of differentiable types (and without an explicit `no_diff`), and propagating them through each instruction in the block. An inst carries a derivative if there one of its operands carries a derivative, and the result type is differentiable. 2. We then assemble a set of instructions that expect a derivative. These are differentiable operands of differentiable functions (unless they have been marked by `no_diff`). We then reverse-propagate this set by adding in all differentiable operands (and repeating this process). 3. During this reverse-propagation, if there is any `OpCall` in the 'expect' set that is not also in the 'produce' set, then we have a situation where the gradient hasn't been explicitly dropped, and we create a user diagnostic. --- ## File: docs/design/autodiff.md Reverse Mode Autodiff (Out of Date) ================================== This document serves as a design reference for reverse-mode auto-diff in the Slang compiler. ## Reverse-Mode Passes Rather than implementing reverse-mode as a separate pass, Slang implements this as a series of independent passes: If a function needs a reverse-mode version generated: - *Linearize* the function, and all dependencies. - *Propagate* differential types through the linearized code. - *Unzip* by moving primal insts to before differential insts. - *Transpose* the differential insts. ## Linearization (Forward-mode) ### Overview (This is a incomplete section. More details coming soon) Consider an arbitrary function `float f(float a, float b, float c, ..., z)` which takes in N inputs and generates one output `y`. Linearization aims to generate the first-order Taylor expansion of f about _all_ of it's inputs. Mathematically, the forward derivative `fwd_f` represents `df/da * (a_0 - a) + df/db * (b_0 - b) + ...`, where `a_0` is the value at which the Taylor expansion was produced. The quantity `a_0 - a` is known as the 'differential' (for brevity we'll denote them da, db, dc, etc..), and there is at-most one differential per input. Thus, the new function's signature should be `fwd_f(float a, float da, float b, float db, float c, float dc, ...)`. For simplicity, we'll use *pairs* instead of interleaving the original and differential parameters. We use the intrinsic `DifferentialPair` (or for short: `DP`) to denote this. The signature we use is then `fwd_f(DP a, DP b, DP c)` An example of linearization: ```C float f(float a, float b) { if (a > 0) { return a + b + 2.0 * a * b; } else { return sqrt(a); } } ``` We'll write out the SSA form of this function. ```C float f_SSA(float a, float b) { bool _b1 = a > 0; if (_b1) { float _t1 = a + b; float _t2 = 2.0 * a; float _t3 = _t2 * b; float _t4 = _t1 + _t3; return _t4; } else { float _t1 = sqrt(a); return _t1; } } DP f_SSA(DP dpa, DP dpb) { bool _b1 = dpa.p > 0; if (_b1) { float _t1 = dpa.p + dpb.p; float _t1_d = dpa.d + dpb.d; float _t2 = 2.0 * dpa.p; float _t2_d = 0.0 * dpa.p + 2.0 * dpa.d; float _t3 = _t2 * dpb.p; float _t3_d = _t2_d * dpb.p + _t2 * dpb.d; float _t4 = _t1 + _t3; float _t4_d = _t1_d + _t3_d; return DP(_t4, _t4_d); } else { DP _t1_dp = sqrt_fwd(dpa); return DP(_t1_dp.p, _t1_dp.d); } } ``` In the result, the primal part of the pair holds the original computation, while the differential part computes the dot product of the differentials with the derivatives of the function's output w.r.t each input. ## Propagation This step takes a linearized function and propagates information about which instructions are computing a differential and which ones are part of the primal (original) computation. Assuming first-order differentiation only: The approach will be to mark any instructions that extract the differential from the differential pair as a differential. Then any instruction that uses the differential is itself marked as a differential and so on. The only exception is the call instruction which is either non-differentiable (do nothing) or differentiable and returns a pair (follow the same process) Here's the above example with propagated type information (we use float.D to denote intermediaries that have been marked as differential, and also expand everything so that each line has a single operation) ```C DP f_SSA_Proped(DP dpa, DP dpb) { bool _b1 = dpa.p > 0; if (_b1) { float _t1 = dpa.p + dpb.p; float.D _q1_d = dpa.d; float.D _q2_d = dpb.d; float.D _t1_d = _q1_d + _q2_d; float _t2 = 2.0 * dpa.p; float.D _q2_d = dpa.d; float.D _q3_d = 2.0 * dpa.d; float _q4 = dpa.p; float.D _q4_d = 0.0 * dpa.p; float.D _t2_d = _q4_d + _q3_d; float _t3 = _t2 * dpb.p; float _q5 = dpb.p; float.D _q6_d = _q5 * _t2_d; float.D _q7_d = dpb.d; float.D _q8_d = _t2 * _q7_d float _t3_d = _q6_d + _q8_d; float _t4 = _t1 + _t3; float.D _t4_d = _t1_d + _t3_d; return DP(_t4, _t4_d); } else { DP _t1_dp = sqrt_fwd(dpa); float _q1 = _t1_dp.p; float.D _q1_d = _t1_dp.d; return DP(_q1, _q1_d); } } ``` ## Unzipping This is a fairly simple process when there is no control flow. We simply move all non-differential instructions to before the first differential instruction. When there is control flow, we need to be a bit more careful: the key is to *replicate* the control flow graph once for primal and once for the differential. Here's the previous example unzipped: ```C DP f_SSA_Proped(DP dpa, DP dpb) { bool _b1 = dpa.p > 0; float _t1, _t2, _q4, _t3, _q5, _t3_d, _t4, _q1; if (_b1) { _t1 = dpa.p + dpb.p; _t2 = 2.0 * dpa.p; _q4 = dpa.p; _t3 = _t2 * dpb.p; _q5 = dpb.p; _t4 = _t1 + _t3; } else { _q1 = sqrt_fwd(DP(dpa.p, 0.0)); } // Note here that we have to 'store' all the intermediaries // _t1, _t2, _q4, _t3, _q5, _t3_d, _t4 and _q1. This is fundamentally // the tradeoff between fwd_mode and rev_mode if (_b1) { float.D _q1_d = dpa.d; float.D _q2_d = dpb.d; float.D _t1_d = _q1_d + _q2_d; float.D _q2_d = dpa.d; float.D _q3_d = 2.0 * dpa.d; float.D _q4_d = 0.0 * dpa.p; float.D _t2_d = _q4_d + _q3_d; float.D _q6_d = _q5 * _t2_d; float.D _q7_d = dpb.d; float.D _q8_d = _t2 * _q7_d float.D _t3_d = _q6_d + _q8_d; float.D _t4_d = _t1_d + _t3_d; return DP(_t4, _t4_d); } else { DP _t1_dp = sqrt_fwd(dpa); float.D _q1_d = _t1_dp.d; return DP(_q1, _q1_d); } } ``` ## Transposition ### Overview This transposition pass _assumes_ that provided function is linear in it's differentials. It is out of scope of this project to attempt to enforce that constraint for user-defined differential code. For transposition we walk all differential instructions in reverse starting from the return statement, and apply the following rules: We'll have an accumulator dictionary `Dictionary accMap` holding assignments for intermediaries which don't have concrete variables. When we add a pair (A, C) and (A, B) already exists, this will form the pair (A, ADD(C, B)) in the dictionary. (ADD will be replaced with a call to `T.dadd` for a generic type T) - If `inst` is a `RETURN(A)`, add pair `(A, d_out)` to `accMap` - If an instruction is `MUL(P, D)` where D is the differential, add pair `(D, MUL(P, accMap[this_inst]))` to `accMap` - If an instruction is `ADD(D1, D2)`, where both D1 and D2 are differentials (this is the only config that should occur), then add pair `(D1, accMap[this_inst])` to `accMap` - If an instruction is `CALL(f_fwd, (P1, D1), (P2, D2), ...)`, create variables D1v, D2v, ... for D1, D2, ..., then replace with `CALL(f_rev, (P1, D1v), (P2, D2v), ..., accMap[this_inst])`, and finally add pairs `(D1, LOAD[D1v]), (D2, LOAD[D2v]), ...` to `accMap` ```C void f_SSA_Rev(inout DP dpa, inout DP dpb, float dout) { bool _b1 = dpa.p > 0; float _t1, _t2, _q4, _t3, _q5, _t3_d, _t4, _q1; if (_b1) { _t1 = dpa.p + dpb.p; _t2 = 2.0 * dpa.p; _q4 = dpa.p; _t3 = _t2 * dpb.p; _q5 = dpb.p; _t4 = _t1 + _t3; } else { _q1 = sqrt_fwd(DP(dpa.p, 0.0)); } // Note here that we have to 'store' all the intermediaries // _t1, _t2, _q4, _t3, _q5, _t3_d, _t4 and _q1. This is fundamentally // the tradeoff between fwd_mode and rev_mode if (_b1) { float.D _t4_rev = d_out; float.D _t1_rev = _t4_rev; float.D _t3_rev = _t4_rev; float.D _q8_rev = _t3_rev; float.D _q6_rev = _t3_rev; float.D _q7_rev = _t2 * _q8_rev; dpb.d += _q7_rev; float.D _t2_rev = _q5 * _q6_rev; float.D _q4_rev = _t2_rev; float.D _q3_rev = _t2_rev; dpa.d += 2.0 * _q3_rev; float.D _q1_rev = _t1_rev; float.D _q2_rev = _t1_rev; dpb.d += _q2_rev; dpa.d += _q1_rev; } else { _q1_rev = d_out; DP dpa_copy; sqrt_rev(dpa_copy, _q1_rev); dpa.d += dpa_copy.d; } } ``` --- ## File: docs/design/capabilities.md Capabilities (Out of Date) ============ Slang aims to be a portable language for shader programming, which introduces two complementary problems: 1. We need a way to indicate that certain constructs (types, functions, etc.) are only allowed on certain targets, so that a user gets a meaningful error if they try to do something that won't work on one or more of the APIs or platforms they want to target. Similarly, the user expects to get an error if they call a fragment-shader-specific function inside of, say, compute shader code, or vice versa. 2. If the same feature can be implemented across multiple platforms, but the best (or only) implementation path differs across platforms, then we need a way to express the platform specific code and pick the right implementation per-target. Item (2) is traditionally handled with preprocessor techniques (e.g., `#ifdef`ing the body of a function based on target platform), but that of course requires that the user invoke the Slang front end once for each target platform, and target-specific coding in a library will then "infect" code that uses that library, forcing them to invoke the front-end once per target as well. We are especially sensitive to this problem in the compiler itself, because we have to author and maintain the Slang standard modules, which needs to (1) expose the capabilities of many platforms and (2) work across all those platforms. It would be very unfortunate if we had to build different copies of our standard modules per-target. The intention in Slang is to solve both of these problems with a system of *capabilities*. What is a capability? --------------------- For our purposes a capability is a discrete feature that a compilation target either does or does not support. We could imagine defining a capability for the presence of texture sampling operations with implicit gradients; this capability would be supported when generating fragment shader kernel code, but not when generating code for other stages. Let's imagine a language syntax that the standard modules could use to define some *atomic* capabilities: ``` capability implicit_gradient_texture_fetches; ``` We can then imagine using attributes to indicate that a function requires a certain capability: ``` struct Texture2D { ... // Implicit-gradient sampling operation. [availableFor(implicit_gradient_texture_fetches)] float4 Sample(SamplerState s, float2 uv); } ``` (Note that the `[availableFor(...)]` syntax is just a straw-man to write up examples, and a better name would be desirable if/when we implement this stuff.) Given those declarations, we could then check when compiling code if the user is trying to call `Texture2D.Sample` in code compiled for a target that *doesn't* support implicit-gradient texture fetches, and issue an appropriate error. The details on how to sequence this all in the compiler will be covered later. Derived Capabilities -------------------- Once we can define atomic capabilities, the next step is to be able to define *derived* capabilities. Let's imagine that we extend our `capability` syntax so that we can define a new capability that automatically implies one or more other capabilities: ``` capability fragment : implicit_gradient_texture_fetches; ``` Here we've said that whenever the `fragment` capability is available, we can safely assume that the `implicit_gradient_texture_fetches` capability is available (but not vice versa). Given even a rudimentary tool like that, we can start to build up capabilities that relate closely to the "profiles" in things like D3D: ``` capability d3d; capability sm_5_0 : d3d; capability sm_5_1 : sm_5_0; capability sm_6_0 : sm_5_1; ... capability d3d11 : d3d, sm_5_0; capability d3d12 : d3d, sm_6_0; capability khronos; capability glsl_400 : khronos; capability glsl_410 : glsl_400; ... capability vulkan : khronos, glsl_450; capability opengl : khronos; ``` Here we are saying that `sm_5_1` supports everything `sm_5_0` supports, and potentially more. We are saying that `d3d12` supports `sm_6_0` but maybe not, e.g., `sm_6_3`. We are expressing that fact that having a `glsl_*` capability means you are on some Khronos API target, but that it doesn't specify which one. (The exact details of these declarations obviously aren't the point; getting a good hierarchy of capabilities will take time.) Capability Composition ---------------------- Sometimes we'll want to give a distinct name to a specific combination of capabilities, but not say that it supports anything new: ``` capability ps_5_1 = sm_5_1 & fragment; ``` Here we are saying that the `ps_5_1` capability is *equivalent* to the combination of `sm_5_1` and `fragment` (that is, if you support both `sm_5_1` and `fragment` then you support `ps_5_1` and vice versa). Compositions should be allowed in `[availableFor(...)]` attributes (e.g., `[availableFor(vulkan & glsl_450)]`), but pre-defined compositions should be favored when possible. When composing things with `&` it is safe for the compiler to filter out redundancies based on what it knows so that, e.g., `ps_5_0 & fragment` resolves to just `ps_5_0`. Once we have an `&` operator for capabilities, it is easy to see that "derived" capabilities are really syntax sugar, so that a derived capability like: ``` capability A : B, C ``` could have been written instead as : ``` capability A_atomic capability A = A_atomic & B & C ``` Where the `A_atomic` capability guarantees that `A` implies `B` and `C` but not vice versa. It is also useful to think of an `|` operator on capabilities. In particular if a function has multiple `[availableFor(...)]` attributes: ``` [availableFor(vulkan & fragment)] [availableFor(d3d12 & fragment)] void myFunc(); ``` This function should be equivalent to one with just a single `[availableFor((vulkan & fragment) | (d3d12 & fragment))]` which is equivalent to `[availableFor((vulkan | d3d12) & fragment)]`. Simplification should generally push toward "disjunctive normal form," though, rather than pursue simplifications like that. Note that we do *not* include negation, so that capabilities are not general Boolean expressions. Validation ---------- For a given function definition `F`, the front end will scan its body and see what it calls, and compose the capabilities required by the called functions using `&` (simplifying along the way). Call the resulting capability (in disjunctive normal form) `R`. If `F` doesn't have an `[availableFor(...)]` attribute, then we can derive its *effective* `[availableFor(...)]` capability as `R` (this probably needs to be expressed as an iterative dataflow problem over the call graph, to handle cycles). If `F` *does* have one or more `[availableFor(...)]` clauses that amount to a declared capability `C` (again in disjunctive normal form), then we can check that `C` implies `R` and error out if it is not the case. A reasonable implementation would track which calls introduced which requirements, and be able to explain *why* `C` does not capture the stated requirements. For a shader entry point, we should check it as if it had an `[availableFor(...)]` that is the OR of all the specified target profiles (e.g., `sm_5_0 | glsl_450 | ...`) ANDed with the specified stage (e.g., `fragment`). Any error here should be reported to the user. If an entry point has an explicit `[availableFor(...)]` then we should AND that onto the profile computed above, so that the user can restrict certain entry points to certain profiles. In order to support separate compilation, the functions that are exported from a module should probably either have explicit availability attributes, or else they will be compiled against a kind of "default capability" used for the whole module. Downstream code that consumes such a module would see declarations with explicit capabilities only. Picking an appropriate "default capability" to use when compiling modules is an important challenge; it would in practice define the "min spec" to use when compiling. Capability Overriding --------------------- It should be possible to define multiple versions of a function, having different `[availableFor(...)]` attributes: ``` [availableFor(vulkan)] void myFunc() { ... } [availableFor(d3d12)] void myFunc() { ... } ``` For front-end checking, these should be treated as if they were a single definition of `myFunc` with an ORed capability (e.g., `vulkan | d3d12`). Overload resolution will pick the "best" candidate at a call site based *only* on the signatures of the function (note that this differs greatly from how profile-specific function overloading works in Cg). The front-end will then generate initial IR code for each definition of `myFunc`. Each of the IR functions will have the *same* mangled name, but different bodies, and each will have appropriate IR decorations to indicate the capabilities it requires. The choice of which definition to use is then put off until IR linking for a particular target. At that point we can look at all the IR functions matching a given mangled name, filter them according to the capabilities of the target, and then select the "best" one. In general a definition `A` of an IR symbol is better than another definition `B` if the capabilities on `A` imply those on `B` but not versa. (In practice this probably needs to be "the capabilities on `A` intersected with those of the target," and similarly for `B`) This approach allows us to defer profile-based choices of functions to very late in the process. The one big "gotcha" to be aware of is when functions are overloaded based on pipeline stage, where we would then have to be careful when generating DXIL or SPIR-V modules with multiple entry points (as a single function `f` might need to be specialized twice if it calls a stage-overloaded function `g`). Capabilities in Other Places ---------------------------- So far I've talked about capabilities on functions, but they should also be allowed on other declarations including: - Types, to indicate that code using that type needs the given capability - Interface conformances, to indicate that a type only conforms to the interface when the capabilities are available - Struct fields, to indicate that the field is only present in the type when the capabilities are present - Extension declarations, to indicate that everything in them requires the specified capabilities We should also provide a way to specify that a `register` or other layout modifier is only applicable for specific targets/stages. Such a capability nominally exists in HLSL today, but it would be much more useful if it could be applied to specify target-API-specific bindings. Only functions should support overloading based on capability. In all other cases there can only be one definition of an entity, and capabilities just decide when it is available. API Extensions as Capabilities ------------------------------ One clear use case for capabilities is to represent optional extensions, including cases where a feature is "built-in" in D3D but requires an extension in Vulkan: ``` capability KHR_secret_sauce : vulkan; [available_for(sm_7_0)] // always available for D3D Shader Model 7.0 [available_for(KHR_secret_sauce)] // Need the "secret sauce" extension for Vulkan void improveShadows(); ``` When generating code for Vulkan, we should be able to tell the user that the `improveShadows()` function requires the given extension. The user should be able to express compositions of capabilities in their `-profile` option (and similarly for the API): ``` slangc code.slang -profile vulkan+KHR_secret_sauce ``` (Note that for the command line, it is beneficial to use `+` instead of `&` to avoid conflicts with shell interpreters) An important question is whether the compiler should automatically infer required extensions without them being specified, so that it produces SPIR-V that requires extensions the user didn't ask for. The argument against such inference is that users should opt in to non-standard capabilities they are using, but it would be unfortunate if this in turn requires verbose command lines when invoking the compiler. It should be possible to indicate the capabilities that a module or entry point should be compiled to use without command-line complications. (A related challenge is when a capability can be provided by two different extensions: how should the compiler select the "right" one to use?) Disjoint Capabilities --------------------- Certain compositions of capabilities make no sense. If a user declared a function as needing `vulkan & d3d12` they should probably get an error message. Knowing that certain capabilities are disjoint can also help improve the overall user experience. If a function requires `(vulkan & extensionA) | (d3d12 & featureb)` and we know we are compiling for `vulkan` we should be able to give the user a pointed error message saying they need to ask for `extensionA`, because adding `featureB` isn't going to do any good. As a first-pass model we could have a notion of `abstract` capabilities that are used to model the root of hierarchies of disjoint capabilities: ``` abstract capability api; abstract capability d3d : api; capability d3d11 : d3d; capability d3d12 : d3d; abstract capability khronos : api; capability vulkan : khronos; capability opengl : khronos; ``` As a straw man: we could have a rule that to decide if non-abstract capabilities `A` and `B` are disjoint, we look for their common ancestor in the tree of capabilities. If the common ancestor is abstract, they are disjoint, and if not they not disjoint. We'd also know that if the user tries to compile for a profile that includes an abstract capability but *not* some concrete capability derived from it, then that is an error (we can't generate code for just `d3d`). The above is an over-simplification because we don't have a *tree* of capabilities, but a full *graph*, so we'd need an approach that works for the full case. Interaction with Generics/Interfaces ------------------------------------ It should be possible for an interface requirement to have a capability requirement attached to it. This would mean that users of the interface can only use the method/type/whatever when the capability is present (just like for any other function): ``` interface ITexture { float4 sampleLevel(float2 uv, float lod); [availableFor(fragment)] float4 sample(float2 uv); // can only call this from fragment code } ``` When implementing an interface, any capability constraints we put on a member that satisfies an interface requirement would need to guarantee that either: - the capabilities on our method are implied by those on the requirement (we don't require more), or - the capabilities on the method are implied by those on the type itself, or its conformance to the interface (you can't use the conformance without the capabilities), or - the capabilities are already implied by those the whole module is being compiled for In each case, you need to be sure that `YourType` can't be passed as a generic argument to some function that uses just the `ITexture` interface above and have them call a method on your type from a profile that doesn't have the required capabilities. Interaction with Heterogeneity ------------------------------ If Slang eventually supports generating CPU code as well as shaders, it should use capabilities to handle the CPU/GPU split similar to how they can be used to separate out vertex- and fragment-shader functionality. Something like a `cpu` profile that works as a catch-all for typical host CPU capabilities would be nice, and could be used as a convenient way to mark "host" functions in a file that is otherwise compiled for a "default profile" that assumes GPU capabilities. Conclusion ---------- Overall, the hope is that in many cases developers will be able to use capability-based partitioning and overloading of APIs to build code that only has to pass through the Slang front-end once, but that can then go through back-end code generation for each target. In cases where this can't be achieved, the way that capability-based overloading is built into the Slang IR design means that we should be able to merge multiple target-specific definitions into one IR module, so that a module can employ target-specific specializations while still presenting a single API to consumers. --- ## File: docs/design/casting.md Casting in the Slang Compiler ============================= The following discussion is about casting within the C++ implementation of the slang compiler. C++'s built in mechanisms for casting (principally dynamic_cast) is problematic within the slang compiler codebase. Code using 'dynamic_cast' requires RTTI information is available, and that a type that uses it must have a vtbl (have at least one virtual member). Some problems with this... * There are types which we want to 'dynamic_cast' that do not have, and we do not want to have a Vtbl (for example Slang::IRInst). * There are types which a 'dynamic_cast' doesn't do quite what we want (for example casting on Type* derived types typically wants to work on their canonical type) * We may want to replace use of dynamic_cast in the future for speed/space or other reasons * It is common in the code base when using a 'smart pointer' type to cast it, but still return a smart pointer To deal with these issues we need casting within Slang to follow it's own methodology. In summary it is as follows... * Use 'as' free function to do a typical 'dynamic like' cast. * 'as' doesn't guarantee the returned pointer points to the same object. * For example with Type* it *actually* does the cast on the canonical type which is often a different object. * If you want to *literally* do a dynamic cast use 'dynamicCast' free function. * This guarantees the returned pointer points to the same object (like normal dynamic_cast) * If you want to return a smart pointer from a cast from a smart pointer use the .as or .dynamicCast *methods* * If you want to determine if an 'as' cast is possible on a smart pointer use the .is method * Doing so will produce more efficient code because a new smart pointer does not need to be constructed These functions will also work with types that do not have Vtbl - like IRInst derived types. Both 'as' and 'dynamicCast' handle the case if the pointer is a nullptr, by returning a nullptr. If the cast succeeds the cast pointer is returned otherwise nullptr is returned. If a cast is performed with a free function it always returns a raw pointer. So why have 'as' and 'dynamicCast' - they seem sort of similar? The primary difference is dynamicCast *must* always return a pointer to the same object, whilst 'as' *can* return a pointer to a different object if that is the desired 'normal' casting behavior for the type. This is the case for Type* when using 'as' it may return a different object - the 'canonical type' for the Type*. For a concrete example take 'NamedExpressionType', its canonical type is the type the name relates to. If you use 'as' on it - it will produce a pointer to a different object, an object that will not be castable back into a NamedExpressionType. Also keep in mind that 'as' behavior is based on the pointer type being cast from. For any pointer to a type derived from Type it will cast the canonical type. **BUT** if the pointer is pointing to a Type derived *object*, but the pointer type is *not* derived from Type (like say RefObject*), then 'as' will behave like dynamicCast. All this being said 'as' in usage is seen as the 'default' way to do a 'dynamic like' cast with these special behaviour appropriate for the type when necessary. By having the free function and method versions of 'as' and 'dynamicCast', you can choose if you want a 'raw' or 'smart' pointer type returned from the cast. If you just want to test if something is a certain type, then using as/dynamicCast free functions is the faster way to do it. If you *know* that a raw pointer is ok, because the object will remain in scope, then again using the free function is better because it does less work. But as the examples following show, care is needed because if you get it wrong the object might go out of scope and leave the raw pointer pointing to a deleted object. When in doubt the safe choice is to typically use .as (or .dynamicCast if appropriate) methods. Following example shows the different types of casting... ```C++ void someFunction(Decl* decl, Type* type) { RefPtr declRefPtr(decl); RefPtr typeRefPtr(type); // Use of as { // Casting with as on a free function returns a raw pointer GenericDecl* genericDeclRaw0 = as(decl); // Free function again returns a raw pointer GenericDecl* genericDeclRaw1 = as(declRefPtr); // Using the as *method* returns a smart pointer holding the cast result RefPtr genericDeclRefPtr0 = declRefPtr.as(); // Of course you can use auto with either auto genericDeclRefPtr1 = declRefPtr.as(); auto genericDeclRaw2 = as(declRefPtr); } // Currently using as on anything not cast *from* Type is the same as dynamicCast. // But on Type* sometimes you may want to control the cast { // With a NamedExpressionType sometimes you don't want 'as' behaviour - if we want to see the information about the name (not the thing // it relates to (the canonical type) NamedExpressionType* namedExpressionRawPtr = dynamicCast(type); // Returns the smart pointer auto namedExpressionRefPtr = typeRefPtr.as(); } ``` It is important to be aware of what style of cast you use where. Take for example the following function ... ```C++ RefPtr substitute(RefPtr expr) const { return DeclRefBase::Substitute(expr); } ``` If you want to do a cast on it, you need to be careful especially about scope, for example... ```C++ RefPtr expr = ...; { // Whoops! This is a problem. When using the free function, the cast is to a *raw* pointer, so obj // receives a raw pointer. When the RefPtr returned from Substitute goes out of scope (when the statement is left) // the ref will be removed and if the ref count was 1 destroyed. Now obj points to a freed object and so a crash is // likely to follow in the future! auto obj = as(substitute(expr)); } // So how do we avoid this? Well it depends what the function is returning and the scope. If it's returning a smart pointer, // you could use the .as method { // This can only compile if it is a smart pointer (raw pointers don't have an as method) auto obj = substitute(expr).as(); } // Another option is to put the created thing in a smart pointer so you know it's in scope { RefPtr sub = substitute(expr); // Ok as long as sub is in scope auto obj = as(sub); } // More awkwardly you could use free function, but assign to a smart pointer, thus maintaining scope { RefPtr obj = as(substitute(expr)); } ``` The following code shows the change in behavior of 'as' is based on the source *pointer* type **NOT** the *object* type.. ```C++ // Derives from Type NamedExpressionType* exprType = ...; // Will be the Type* of the *canonical* type, because the pointer is Type derived and we are using as! Type* type0 = as(exprType); // It' going to be pointing to a different object, because type0 is the cast of the *canonical* type, because exprType derives from Type SLANG_ASSERT(type0 != exprType); // If I do a dynamicCast the result is either nullptr or a pointer that *must* point to the same object Type* type1 = dynamicCast(exprType); SLANG_ASSERT(type1 == exprType); // Here, the pointer is pointing to a NamedExpressionType derived object. Which derives from Type. BUT our pointer here does *not* derive from type. RefObject* refObj = exprType; // 'as' just looks at the from type, and it doesn't derive from Type (it's just RefObject), so it does regular as, which is dynamicCast Type* type2 = as(refObject); SLANG_ASSERT(type2 == exprType); // Finally... // Is true even though exprType is a NamedExpression, because the cast is on the canonical type SLANG_ASSERT(as(exprType) == nullptr); // dynamicCast is always the same object returned, so must match SLANG_ASSERT(dynamicCast(exprType) == exprType); ``` --- ## File: docs/design/coding-conventions.md Slang Project Coding Conventions ================================ Principles ---------- This document attempts to establish conventions to be used in the Slang codebase. We have two goals for this convention. The first goal is to make the code look relatively consistent so that it is easy to navigate and understand for contributors. Having varying styles across different modules, files, functions, or lines of code makes the overall design and intention of the codebase harder to follow. The second goal is to minimize the scope complexity of diffs when multiple maintainers work together on the codebase. In the absence of an enforced style, developers tend to "clean up" code they encounter to match their personal preferences, and in so doing create additional diffs that increase the chances of merge conflicts and pain down the line. Because the Slang codebase has passed through many hands and evolved without a pre-existing convention, these two goals can come into conflict. We encourage developers to err on the side of leaving well enough alone (favoring the second goal). Don't rewrite or refactor code to match these conventions unless you were already going to have to touch all of those lines of code anyway. Note that external code that is incorporated into the project is excluded from all of these conventions. Languages --------- ### C++ Most code in the Slang project is implemented in C++. We currently assume support for some C++11 idioms, but have explicitly avoided adding dependencies on later versions. As a general rule, be skeptical of "modern C++" ideas unless they are clearly better to simpler alternatives. We are not quite in the realm of "Orthodox C++", but some of the same guidelines apply: * Don't use exceptions for non-fatal errors (and even then support a build flag to opt out of exceptions) * Don't use the built-in C++ RTTI system (home-grown is okay) * Don't use the C++ variants of C headers (e.g., `` instead of ``) * Don't use the STL containers * Don't use iostreams The compiler implementation does not follow some of these guidelines at present; that should not be taken as an excuse to further the proliferation of stuff like `dynamic_cast`. Do as we say, not as we do. Some relatively recent C++ features that are okay to use: * Rvalue references for "move semantics," but only if you are implementing performance-critical containers or other code where this really matters. * `auto` on local variables, if the expected type is clear in context * Lambdas are allowed, but think carefully about whether just declaring a subroutine would also work. * Using `>>` to close multiple levels of templates, instead of `> >` (but did you really need all those templates?) * `nullptr` * `enum class` * Range-based `for` loops * `override` * Default member initializers in `class`/`struct` bodies Templates are suitable in cases where they improve clarity and type safety. As a general rule, it is best when templated code is kept minimal, and forwards to a non-templated function that does the real work, to avoid code bloat. Any use of template metaprogramming would need to prove itself exceptionally useful to pay for the increase in cognitive complexity. We don't want to be in the business of maintaining "clever" code. As a general rule, `const` should be used sparingly and only with things that are logically "value types." If you find yourself having to `const`-qualify a lot of member function in type that you expect to be used as a heap-allocated object, then something has probably gone wrong. As a general rule, default to making the implementation of a type `public`, and only encapsulate state or operations with `private` when you find that there are complex semantics or invariants that can't be provided without a heavier hand. ### Slang The Slang project codebase also includes `.slang` files implementing the Slang core module, as well as various test cases and examples. The conventions described here are thus the "official" recommendations for how users should format Slang code. To the extent possible, we will try to apply the same basic conventions to both C++ and Slang. In places where we decide that the two languages merit different rules, we will point it out. Files and Includes ------------------ ### File Names All files and directories that are added to codebase should have names that contain only ASCII lower-case letters, digits, dots (`.`) and dashes (`-`). Operating systems still vary greatly in their handling of case sensitivity for file names, and non-ASCII code points are handled with even less consistency; sticking to a restricted subset of ASCII helps avoids some messy interactions between case-insensitive file systems and case-sensitive source-control systems like Git. As with all these conventions, files from external projects are exempted from these restrictions. ### Naming of Source and Header Files In general the C++ codebase should be organized around logical features/modules/subsystem, each of which has a single `.h` file and zero or more `.cpp` files to implement it. If there is a single `.cpp` file, its name should match the header: e.g., `parser.h` and `parser.cpp`. If there is more than one `.cpp` file, their names should start with the header name: e.g., `parser.h` and `parser-decls.cpp` and `parser-exprs.cpp`. If there are declarations that need to be shared by the `.cpp` files, but shouldn't appear in the public interface, then can go in a `*-impl.h` header (e.g., `parser-impl.h`). Use best judgement when deciding what counts as a "feature." One class per file is almost always overkill, but the codebase currently leans too far in the other direction, with some oversized source files. ### Headers Every header file should have an include guard. Within the implementation we can use `#pragma once`, but exported API headers (`slang.h`) should use traditional `#ifdef` style guards (and they should be consumable as both C and C++). A header should include or forward-declare everything it needs in order to compile. It is *not* up to the programmer who `#include`s a header to sort out the dependencies. Avoid umbrella or "catch-all" headers. ### Source Files Every source file should start by including the header for its feature/module, before any other includes (this helps ensure that the header correctly includes its dependencies). Functions that are only needed within that one source file can be marked `static`, but we should avoid using the same name for functions in different files (in order to support lumped/unified builds). ### Includes Cross-directory project-internal includes should name a header from an include root and should not use `../` to traverse directories. The `source/` directory is an include root, so cross-module includes should use paths such as: ```c++ #include "compiler-core/slang-source-loc.h" #include "core/slang-string.h" ``` Targets that consume headers under another project directory, such as `tools/` or `examples/`, should expose that directory through their CMake include directories and use paths relative to that root. Prefer a private include directory unless it is part of the target's public interface. In general, includes should be grouped as follows: * First, the correspodning feature/module header, if we are in a source file * Next, any `<>`-enlosed includes for system/OS headers * Next, any `""`-enclosed includes for external/third-part code that is stored in the project repository * Finally, any includes for other features in the project Within each group, includes should be sorted alphabetically. If this breaks because of ordering issues for system/OS/third-party headers (e.g., `` must be included before ``), then ideally those includes should be mediated by a Slang-project-internal header that features can include. Namespaces ---------- Favor fewer namespaces when possible. Small programs may not need any. All standard module code that a Slang user might link against should go in the `Slang` namespace for now, to avoid any possibility of clashes in a static linking scenario. The public C API is obviously an exception to this. Code Formatting ------------------------------ - For C++ files, please format using `clang-format`; `.clang-format` files in the source tree define the style. - For CMake files, please format using `gersemi` - For shell scripts, please format using `shfmt` - For YAML files, please use `prettier` The formatting for the codebase is overall specified by the [`extras/formatting.sh`](./extras/formatting.sh) script. If you open a pull request and the formatting is incorrect, you can comment `/format` and a bot will format your code for you. Naming ------ ### Casing Types should in general use `UpperCamelCase`. This includes `struct`s, `class`es, `enum`s and `typedef`s. Values should in general use `lowerCamelCase`. This includes functions, methods, local variables, global variables, parameters, fields, etc. Macros should in general use `SCREAMING_SNAKE_CASE`. It is important to prefix all macros (e.g., with `SLANG_`) to avoid collisions, since `namespace`s don't affect macros). In names using camel case, acronyms and initialisms should appear eniterly in either upper or lower case (e.g., `D3DThing d3dThing`) and not be capitalized as if they were ordinary words (e.g., `D3dThing d3dThing`). Note that this also applies to uses of "ID" as an abbreviation for "identifier" (e.g., use `nodeID` instead of `nodeId`). ### Prefixes Prefixes based on types (e.g., `p` for pointers) should never be used. Global variables should have a `g` prefix, e.g. `gCounter`. Non-`const` `static` class members can have an `s` prefix if that suits your fancy. Of course, both of these should be avoided, so this shouldn't come up often. Constant data (in the sense of `static const`) should have a `k` prefix. In contexts where "information hiding" is relevant/important, such as when a type has both `public` and `private` members, or just has certain operations/fields that are considered "implementation details" that most clients should not be using, an `m_` prefix on member variables and a `_` prefix on member functions is allowed (but not required). In function parameter lists, an `in`, `out`, or `io` prefix can be added to a parameter name to indicate whether a pointer/reference/buffer is intended to be used for input, output, or both input and output. For example: ```c++ void copyData(void* outBuffer, void const* inBuffer, size_t size); Result lookupThing(Key k, Thing& outThing); void maybeAppendExtraNames(std::vector& ioNames); ``` Public C APIs will prefix all symbol names while following the casing convention (e.g. `SlangModule`, `slangLoadModule`, etc.). ### Internal `spirv_asm` result registers Result registers named inside a `spirv_asm` block in the core module (the `*.meta.slang` files) must be given a `__` prefix, e.g. `%__result` rather than `%result`. The SPIR-V emitter emits an `OpName` for every named `spirv_asm` register, so an un-prefixed internal name such as `%result` surfaces in the generated module's debug names where it can be mistaken for a user symbol. The `__` prefix marks these names as compiler-internal, matching the convention above for other not-for-normal-use core-module identifiers. This is enforced at parse time: the parser asserts that every named `spirv_asm` register in core-module code begins with `__`. User-authored `spirv_asm` blocks are unaffected — they may name registers however they like. ### Enums C-style `enum` should use the following convention: ```c++ enum Color { kColor_Red, kColor_Green, kColor_Blue, kColorCount, }; ``` When using `enum class`, drop the `k` and type name as prefix, but retain the `UpperCamelCase` tag names: ```c++ enum class Color { Red, Green, Blue, Count, }; ``` When defining a set of flags, separate the type definition from the `enum`: ```c++ typedef unsigned int Axes; enum { kAxes_None = 0, kAxis_X = 1 << 0, kAxis_Y = 1 << 1, kAxis_Z = 1 << 2, kAxes_All = kAxis_X | kAxis_Y | kAxis_Z, }; ``` Note that the type name reflects the plural case, while the cases that represent individual bits are named with a singular prefix. In public APIs, all `enum`s should use the style of separating the type definition from the `enum`, and all cases should use `SCREAMING_SNAKE_CASE`: ```c++ typedef unsigned int SlangAxes; enum { SLANG_AXES_NONE = 0, SLANG_AXIS_X = 1 << 0, SLANG_AXIS_Y = 1 << 1, SLANG_AXIS_Z = 1 << 2, SLANG_AXES_ALL = SLANG_AXIS_X | SLANG_AXIS_Y | SLANG_AXIS_Z, }; ``` ### General Names should default to the English language and US spellings, to match the dominant conventions of contemporary open-source projects. Function names should either be named with action verbs (`get`, `set`, `create`, `emit`, `parse`, etc.) or read as questions (`isEnabled`, `shouldEmit`, etc.). Whenever possible, compiler concepts should be named using the most widely-understood term available: e.g., we use `Token` over `Lexeme`, and `Lexer` over `Scanner` simply because they appear to be the more common names. Avoid abbreviations and initialisms unless they are already widely established across the codebase; a longer name may be cumbersome to write in the moment, but the code will probably be read many more times than it is written, so clarity should be preferred. An important exception to this is common compiler concepts or techniques which may have laboriously long names: e.g., Static Single Assignment (SSA), Sparse Conditional Copy Propagation (SCCP), etc. One gotcha particular to compiler front-ends is that almost every synonym for "type" has some kind of established technical meaning; most notably the term "kind" has a precise meaning that is relevant in our domain. It is common practice in C and C++ to define tagged union types with a selector field called a "type" or "kind," which does not usually match this technical definition. If a developer wants to avoid confusion, they are encouraged to use the term "flavor" instead of "type" or "kind" since this term (while a bit silly) is less commonly used in the literature. Comments and Documentation -------------------------- You probably know the drill: comments are good, but an out-of-date comment can be worse than no comment at all. Try to write comments that explain the "why" of your code more than the "what." When implementing a textbook algorithm or technique, it may help to imagine giving the reviewer of your code a brief tutorial on the topic. In cases where comments would benefit from formatting, use Markdown syntax. We do not currently have a setup for extracting documentation from comments, but if we add one we will ensure that it works with Markdown. When writing comments, please be aware that your words could be read by many people, from a variety of cultures and backgrounds. Default to a plain-spoken and professional tone and avoid using slang, idiom, profanity, etc. ## 2. Official Technical Reference & Guides (shader-slang/shader-slang.github.io) ## File: README.md # Shader-Slang.com website This theme is losely based on the Serif theme for Jekyll by [Zerostatic Themes](https://www.zerostatic.io). ## Installation ### Installing Ruby & Jekyll If this is your first time using Jekyll, please follow the [Jekyll docs](https://jekyllrb.com/docs/installation/) and make sure your local environment (including Ruby) is setup correctly. ### Installing Theme Download or clone the theme. To run the theme locally, navigate to the theme directory and run: ``` bundle install ``` To start the Jekyll local development server. ``` bundle exec jekyll serve ``` To build the theme. ``` bundle exec jekyll build ``` ### Github Pages This theme has been tested to work with Github Pages (and Github Project Pages). When using Github Pages you will need to update the `baseurl` in the `_config.yml` otherwise all the css, images and paths will be broken. For example the site https://shader-slang.com would have `baseurl: "/"` --- ## File: docs/index.md --- title: Documentation layout: docs permalink: "/docs/" ---

Overview

{% assign overview_items = site.data.documentation.overview %} {% for item in overview_items %} {% if item.title and item.title != '' %}

{{ item.title }}

{{ item.description }} {% if item.link_url and item.link_url != '' %} [{{ item.link_label }}]({{ item.link_url }}) {% else %} {{ item.link_label }} {% endif %} {% endif %} {% endfor %}

Articles

{% assign articles_items = site.data.documentation.articles %} {% for item in articles_items %} {% if item.title and item.title != '' %}

{{ item.title }}

{{ item.description }} {% if item.link_url and item.link_url != '' %} [{{ item.link_label }}]({{ item.link_url }}) {% else %} {{ item.link_label }} {% endif %} {% endif %} {% endfor %}

Tutorials

{% assign tutorials_items = site.data.documentation.tutorials %} {% for item in tutorials_items %} {% if item.title and item.title != '' %}

{{ item.title }}

{{ item.description }} {% if item.link_url and item.link_url != '' %} [{{ item.link_label }}]({{ item.link_url }}) {% else %} {{ item.link_label }} {% endif %} {% endif %} {% endfor %}

Contributions

If you’d like to contribute to the project, we are excited to have your input. Our community pages provides information on our structure and our process for accepting contributions. Making Community Contributions --- ## File: docs/getting-started.md --- title: Getting Started layout: getting-started permalink: "/docs/getting-started/" intro_image_absolute: true intro_image_hide_on_mobile: false ---

Try Slang

Try slang locally in your browser without downloading or installing anything with the Slang Playground. Try the Slang Playground

Machine Learning with Slang

The SlangPy package seamlessly bridges the graphics and machine learning ecosystems in Python. You can write Slang code to implement complex data structures and algorithms, or to use rasterization and ray-tracing. Slang code can be integrated directly into your Python/PyTorch training loop with SlangPy. No boilerplate, no C++. Try SlangPy

Learn How to Become a Slang Expert

Slang code is highly portable, but can still leverage unique platform capabilities, including the latest features in Direct3D and Vulkan. For example, developers can make full use of [pointers](https://docs.shader-slang.org/en/latest/external/slang/docs/user-guide/03-convenience-features.html#pointers-limited) when generating SPIR-V. Slang's capability system helps applications manage feature set differences across target platforms by ensuring code only uses available features during the type-checking step, before generating final code. Additionally, Slang provides flexible interop features to enable directly embedding target code or SPIR-V into generated shaders. {% assign gs_items = site.data.getting-started.items %} {% for item in gs_items %} {% if item.title and item.title != '' %} {{ item.icon }}

{{ item.title }}

{{ item.description }} {% if item.link_url and item.link_url != '' %} [{{ item.link_label }} >]({{ item.link_url }}) {% else %} {{ item.link_label }} {% endif %} {% endif %} {% endfor %}