Added [AutoPyBindCUDA] for automatic kernel binding + [PyExport] for exporting type information (shader-slang#3209)

* Initial: add a DiffTensor impl

* Auto-binding and diff tensor implementations now work

* Refactored diff-tensor implementation + added py-export for struct types

* Cleanup

* Update slang-ir-pytorch-cpp-binding.cpp

* Updated test names

* Update autodiff-data-flow.slang.expected

* Add more versions of load/store & default generic args for DiffTensorView.

* Add diagnostic for default generic arg and more tests

* Add more `[AutoPyBind]` tests

Author: saipraveenb25

source/slang/core.meta.slang

@@ -2263,6 +2263,12 @@ attribute_syntax [CudaHost] : CudaHostAttribute;
 __attributeTarget(FuncDecl)
 attribute_syntax [CudaKernel] : CudaKernelAttribute;
 
+__attributeTarget(FuncDecl)
+attribute_syntax[AutoPyBindCUDA] : AutoPyBindCudaAttribute;
+
+__attributeTarget(AggTypeDecl)
+attribute_syntax [PyExport(name: String)] : PyExportAttribute;
+
 __attributeTarget(InterfaceDecl)
 attribute_syntax [COM(guid: String)] : ComInterfaceAttribute;
 

source/slang/diff.meta.slang

@@ -264,6 +264,298 @@ extension TensorView<float>
     void InterlockedCompareExchange(vector<uint, N> index, float compare, float val);
 }
 
+interface IDiffTensorWrapper
+{
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward(uint offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward_2(uint2 offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward_3(uint3 offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward_4(uint4 offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward(uint offset, T dOut);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward_2(uint2 offset, T dOut);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward_3(uint3 offset, T dOut);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward_4(uint4 offset, T dOut);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward(uint offset, T dx);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward_2(uint2 offset, T dx);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward_3(uint3 offset, T dx);
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward_4(uint4 offset, T dx);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward(uint offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward_2(uint2 offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward_3(uint3 offset);
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward_4(uint4 offset);
+};
+
+struct AtomicAdd : IDiffTensorWrapper
+{
+    TensorView<float> diff;
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward(uint i)
+    {
+        return __realCast<T, float>(diff.load(i));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward_2(uint2 i)
+    {
+        return __realCast<T, float>(diff.load(i.x, i.y));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward_3(uint3 i)
+    {
+        return __realCast<T, float>(diff.load(i.x, i.y, i.z));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T load_forward_4(uint4 i)
+    {
+        return __realCast<T, float>(diff.load(i.x, i.y, i.z, i.w));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward(uint i, T dOut)
+    {
+        float oldVal;
+        diff.InterlockedAdd(i, __realCast<float, T>(dOut), oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward_2(uint2 i, T dOut)
+    {
+        float oldVal;
+        diff.InterlockedAdd(i, __realCast<float, T>(dOut), oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward_3(uint3 i, T dOut)
+    {
+        float oldVal;
+        diff.InterlockedAdd(i, __realCast<float, T>(dOut), oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void load_backward_4(uint4 i, T dOut)
+    {
+        float oldVal;
+        diff.InterlockedAdd(i, __realCast<float, T>(dOut), oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward(uint i, T dx)
+    {
+        diff.store(i, __realCast<float, T>(dx));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward_2(uint2 i, T dx)
+    {
+        diff.store(i.x, i.y, __realCast<float, T>(dx));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward_3(uint3 i, T dx)
+    {
+        diff.store(i.x, i.y, i.z, __realCast<float, T>(dx));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    void store_forward_4(uint4 i, T dx)
+    {
+        diff.store(i.x, i.y, i.z, i.w, __realCast<float, T>(dx));
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward(uint i)
+    {
+        float oldVal;
+        diff.InterlockedExchange(i, (float)0, oldVal);
+        return __realCast<T, float>(oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward_2(uint2 i)
+    {
+        float oldVal;
+        diff.InterlockedExchange(i, (float)0, oldVal);
+        return __realCast<T, float>(oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward_3(uint3 i)
+    {
+        float oldVal;
+        diff.InterlockedExchange(i, (float)0, oldVal);
+        return __realCast<T, float>(oldVal);
+    }
+
+    __generic<T : __BuiltinFloatingPointType>
+    T store_backward_4(uint4 i)
+    {
+        float oldVal;
+        diff.InterlockedExchange(i, (float)0, oldVal);
+        return __realCast<T, float>(oldVal);
+    }
+};
+
+__generic<T: __BuiltinFloatingPointType = float, A : IDiffTensorWrapper = AtomicAdd>
+struct DiffTensorView
+{
+    TensorView<T> primal;
+    A diff;
+
+    uint size(uint i)
+    {
+        return primal.size(i);
+    }
+
+    [BackwardDerivative(load_backward)]
+    [ForwardDerivative(load_forward)]
+    T load(uint i) { return primal.load(i); }
+
+    [BackwardDerivative(load_backward)]
+    [ForwardDerivative(load_forward)]
+    T load(uint2 i) { return primal.load(i.x, i.y); }
+
+    [BackwardDerivative(load_backward)]
+    [ForwardDerivative(load_forward)]
+    T load(uint3 i) { return primal.load(i.x, i.y, i.z); }
+
+    [BackwardDerivative(load_backward)]
+    [ForwardDerivative(load_forward)]
+    T load(uint4 i) { return primal.load(i.x, i.y, i.z, i.w); }
+
+    DifferentialPair<T> load_forward(uint x)
+    {
+        return diffPair(primal.load(x), reinterpret<T.Differential, T>(diff.load_forward<T>(x)));
+    }
+
+    DifferentialPair<T> load_forward(uint2 x)
+    {
+        return diffPair(primal.load(x.x, x.y), reinterpret<T.Differential, T>(diff.load_forward_2<T>(x)));
+    }
+
+    DifferentialPair<T> load_forward(uint3 x)
+    {
+        return diffPair(primal.load(x.x, x.y, x.z), reinterpret<T.Differential, T>(diff.load_forward_3<T>(x)));
+    }
+
+    DifferentialPair<T> load_forward(uint4 x)
+    {
+        return diffPair(primal.load(x.x, x.y, x.z, x.w), reinterpret<T.Differential, T>(diff.load_forward_4<T>(x)));
+    }
+
+    void load_backward(uint x, T.Differential dOut)
+    {
+        diff.load_backward<T>(x, reinterpret<T, T.Differential>(dOut));
+    }
+
+    void load_backward(uint2 x, T.Differential dOut)
+    {
+        diff.load_backward_2<T>(x, reinterpret<T, T.Differential>(dOut));
+    }
+
+    void load_backward(uint3 x, T.Differential dOut)
+    {
+        diff.load_backward_3<T>(x, reinterpret<T, T.Differential>(dOut));
+    }
+
+    void load_backward(uint4 x, T.Differential dOut)
+    {
+        diff.load_backward_4<T>(x, reinterpret<T, T.Differential>(dOut));
+    }
+
+    [BackwardDerivative(store_backward)]
+    [ForwardDerivative(store_forward)]
+    void store(uint x, T val) { primal.store(x, val); }
+
+    [BackwardDerivative(store_backward)]
+    [ForwardDerivative(store_forward)]
+    void store(uint2 x, T val) { primal.store(x.x, x.y, val); }
+
+    [BackwardDerivative(store_backward)]
+    [ForwardDerivative(store_forward)]
+    void store(uint3 x, T val) { primal.store(x.x, x.y, x.z, val); }
+
+    [BackwardDerivative(store_backward)]
+    [ForwardDerivative(store_forward)]
+    void store(uint4 x, T val) { primal.store(x.x, x.y, x.z, x.w, val); }
+
+    void store_forward(uint x, DifferentialPair<T> dpval)
+    {
+        primal.store(x, dpval.p);
+        diff.store_forward<T>(x, reinterpret<T, T.Differential>(dpval.d));
+    }
+
+    void store_forward(uint2 x, DifferentialPair<T> dpval)
+    {
+        primal.store(x.x, x.y, dpval.p);
+        diff.store_forward_2<T>(x, reinterpret<T, T.Differential>(dpval.d));
+    }
+
+    void store_forward(uint3 x, DifferentialPair<T> dpval)
+    {
+        primal.store(x.x, x.y, x.z, dpval.p);
+        diff.store_forward_3<T>(x, reinterpret<T, T.Differential>(dpval.d));
+    }
+
+    void store_forward(uint4 x, DifferentialPair<T> dpval)
+    {
+        primal.store(x.x, x.y, x.z, x.w, dpval.p);
+        diff.store_forward_4<T>(x, reinterpret<T, T.Differential>(dpval.d));
+    }
+
+    void store_backward(uint x, inout DifferentialPair<T> dpval)
+    {
+        dpval = diffPair(dpval.p, reinterpret<T.Differential, T>(diff.store_backward<T>(x)));
+    }
+
+    void store_backward(uint2 x, inout DifferentialPair<T> dpval)
+    {
+        dpval = diffPair(dpval.p, reinterpret<T.Differential, T>(diff.store_backward_2<T>(x)));
+    }
+
+    void store_backward(uint3 x, inout DifferentialPair<T> dpval)
+    {
+        dpval = diffPair(dpval.p, reinterpret<T.Differential, T>(diff.store_backward_3<T>(x)));
+    }
+
+    void store_backward(uint4 x, inout DifferentialPair<T> dpval)
+    {
+        dpval = diffPair(dpval.p, reinterpret<T.Differential, T>(diff.store_backward_4<T>(x)));
+    }
+};
+
 /// Represents the handle of a Torch tensor object.
 __generic<T>
 __intrinsic_type($(kIROp_TorchTensorType))

source/slang/slang-ast-modifier.h

@@ -1134,6 +1134,18 @@ class CudaHostAttribute : public Attribute
     SLANG_AST_CLASS(CudaHostAttribute)
 };
 
+class AutoPyBindCudaAttribute : public Attribute
+{
+    SLANG_AST_CLASS(AutoPyBindCudaAttribute)
+};
+
+class PyExportAttribute : public Attribute
+{
+    SLANG_AST_CLASS(PyExportAttribute)
+
+    String name;
+};
+
 class PreferRecomputeAttribute : public Attribute
 {
     SLANG_AST_CLASS(PreferRecomputeAttribute)

source/slang/slang-check-modifier.cpp

@@ -757,6 +757,19 @@ namespace Slang
 
             knownBuiltinAttr->name = name;
         }
+        else if (auto pyExportAttr = as<PyExportAttribute>(attr))
+        {
+            // Check name string.
+            SLANG_ASSERT(attr->args.getCount() == 1);
+
+            String name;
+            if(!checkLiteralStringVal(attr->args[0], &name))
+            {
+                return false;
+            }
+
+            pyExportAttr->name = name;
+        }
         else
         {
             if(attr->args.getCount() == 0)