Limitations of automatic differentiation

[chenzhekl]

Thanks for releasing this great shader language! I would like to use slang for GPGPU tasks but have a few questions:

1. I read the guide of integrating slang with PyTorch. It mentioned that the 3x3 box filter cannot be automatically differentiated without introducing a custom backward function. I am wondering what the limitation is here. Does slang only support autodiff functions whose output and input are 1:1 mapped? Would that imply for functions with, for example, 3 differentiable outputs and 2 differentiable inputs, it is better to write the backward pass completely by ourselves?
2. It is very convenient to write a compute shader that specifically targets CUDA and PyTorch with the help of slangpy. Are there any approaches to simplifying code reuse if I would like to deploy the same algorithm to the Vulkan compute backend?

It would be greatly appreciated if you could help!

[csyonghe]

Thank you for your interest in trying out Slang for your project! The tutorial might be a little misleading in that Slang does support auto-diffing the box filter. It is just that Slang cannot differentiate through a global memory access operation. In the case where you are accessing a global buffer/texture, you need to wrap that with a function, e.g. the getInputElement function in the example, and through the custom derivative of that function, specify how you want the gradients to be accumulated in the target buffer. In the example, we simply call InterlockedAdd to natively accumulate the value from all threads into a single buffer location.

This certainly leads to high contention (and therefore low performance), and this is exactly the reason why Slang does not do it automatically for you. By providing a custom derivative to getInputElement, you are given the opportunity to implementing much faster accumulating techniques than naive atomicAdd.

If you wish to use Slang code in Vulkan, just write a normal compute shader entry point can call your existing Slang functions. You can then compile your compute shader into glsl / spirv and use it in Vulkan. You can use the slangc tool to do so:

slangc myShader.slang -entry computeMain -stage compute -target spirv

[chenzhekl]

Hi Yong,

Thanks so much for your detailed explanation! It makes sense that a global access requires special care for correct gradient accumulation. May I have one more question? Regarding point 2, I find it hard to write code that works with both TensorView<T> and StructuredBuffer<T>. Would you mind providing some advice?

Thanks!

[csyonghe]

I see. We are working on a new language feature that would allow you to define generic methods over different container types. For example, with #3281, you can define:

interface IBuffer<T> {
   T get(int index);
   int getCount();
}

And you can use extensions to make both StructuredBuffer and TensorView conform to IBuffer:

__generic<T> extension TensorView<T>: IBuffer<T>
{
    T get(int index) { return this[index]; }
    int getCount() { return dim(0); }
}

// Define a new wrapper type for StructuredBuffer because StructuredBuffer itself doesn't contain count.
struct MyBuffer<T> : IBuffer<T> {
  int count;
  StructuredBuffer<T> buffer;
  T get(int index) { return buffer[index]; }
  int getCount() { return count; }
}

Then you can write your generic function as:

float myFunc(IBuffer<float> buffer) {}

That can be called with either a MyBuffer or a TensorView.

[chenzhekl]

Thank you so much for your advice! #3281 is indeed what we need.