Added `[AutoPyBindCUDA]` for automatic kernel binding + `[PyExport]` for exporting type information (#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

17 files changed, 224 insertions, 13 deletions

diff --git a/tests/autodiff/autopybind-basic.slang b/tests/autodiff/autopybind-basic.slang
new file mode 100644
index 000000000..5d4409474
--- /dev/null
+++ b/tests/autodiff/autopybind-basic.slang
@@ -0,0 +1,25 @@
+//TEST:SIMPLE(filecheck=CUDA): -target cuda -line-directive-mode none
+//TEST:SIMPLE(filecheck=TORCH): -target torch -line-directive-mode none
+
+// CUDA: __global__ void __kernel__myKernel(TensorView inValues_[[#]], TensorView outValues_[[#]])
+[AutoPyBindCUDA]
+[CudaKernel]
+void myKernel(TensorView<float> inValues, TensorView<float> outValues)
+{
+    if (cudaThreadIdx().x > 0)
+        return;
+    outValues.store(cudaThreadIdx().x, sin(inValues.load(cudaThreadIdx().x)));
+}
+
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void __kernel__myKernel(TensorView {{[[:alnum:]_]+}}, TensorView {{[[:alnum:]_]+}});
+
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void myKernel(std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, torch::Tensor {{[[:alnum:]_]+}}, torch::Tensor {{[[:alnum:]_]+}})
+
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: std::tuple<std::tuple<const char*, const char*, const char*, const char*>, std::tuple<const char*, const char*>, const char*, const char*> __funcinfo__myKernel()
+
+// TORCH:      m.def("myKernel", &myKernel, "myKernel");
+
+// TORCH:      m.def("__funcinfo__myKernel", &__funcinfo__myKernel, "__funcinfo__myKernel");
\ No newline at end of file
diff --git a/tests/autodiff/autopybind-differentiable.slang b/tests/autodiff/autopybind-differentiable.slang
new file mode 100644
index 000000000..9595a09d2
--- /dev/null
+++ b/tests/autodiff/autopybind-differentiable.slang
@@ -0,0 +1,44 @@
+//TEST:SIMPLE(filecheck=CUDA): -target cuda -line-directive-mode none
+//TEST:SIMPLE(filecheck=TORCH): -target torch -line-directive-mode none
+
+// CUDA: __global__ void __kernel__myKernel_bwd_diff(DiffTensorView_[[#]] {{[[:alnum:]_]+}}, DiffTensorView_[[#]] {{[[:alnum:]_]+}})
+// CUDA: __global__ void __kernel__myKernel_fwd_diff(DiffTensorView_[[#]] {{[[:alnum:]_]+}}, DiffTensorView_[[#]] {{[[:alnum:]_]+}})
+// CUDA: __global__ void __kernel__myKernel(DiffTensorView_[[#]] {{[[:alnum:]_]+}}, DiffTensorView_[[#]] {{[[:alnum:]_]+}})
+
+[AutoPyBindCUDA]
+[Differentiable]
+[CudaKernel]
+void myKernel(DiffTensorView inValues, DiffTensorView outValues)
+{
+    if (cudaThreadIdx().x > 0)
+        return;
+    outValues.store(cudaThreadIdx().x, sin(inValues.load(cudaThreadIdx().x)));
+}
+
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void __kernel__myKernel_bwd_diff(DiffTensorView_[[#]] {{[[:alnum:]_]+}}, DiffTensorView_[[#]] {{[[:alnum:]_]+}})
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void __kernel__myKernel_fwd_diff(DiffTensorView_[[#]] {{[[:alnum:]_]+}}, DiffTensorView_[[#]] {{[[:alnum:]_]+}})
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void __kernel__myKernel(DiffTensorView_[[#]] {{[[:alnum:]_]+}}, DiffTensorView_[[#]] {{[[:alnum:]_]+}})
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void myKernel(std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, std::tuple<torch::Tensor>> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, std::tuple<torch::Tensor>> {{[[:alnum:]_]+}})
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: std::tuple<std::tuple<const char*, const char*, const char*, const char*>, std::tuple<const char*, const char*>, const char*, const char*> __funcinfo__myKernel()
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void myKernel_fwd_diff(std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, std::tuple<torch::Tensor>> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, std::tuple<torch::Tensor>> {{[[:alnum:]_]+}})
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void myKernel_bwd_diff(std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, std::tuple<torch::Tensor>> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, std::tuple<torch::Tensor>> {{[[:alnum:]_]+}})
+//
+// TORCH:     {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: std::tuple<std::tuple<const char*, const char*>, std::tuple<const char*, const char*>> __typeinfo__DiffTensorView()
+//
+// TORCH:     {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: std::tuple<std::tuple<const char*>, std::tuple<const char*>> __typeinfo__AtomicAdd()
+//
\ No newline at end of file
diff --git a/tests/autodiff/autopybind-struct.slang b/tests/autodiff/autopybind-struct.slang
new file mode 100644
index 000000000..d94218494
--- /dev/null
+++ b/tests/autodiff/autopybind-struct.slang
@@ -0,0 +1,35 @@
+//TEST:SIMPLE(filecheck=CUDA): -target cuda -line-directive-mode none
+//TEST:SIMPLE(filecheck=TORCH): -target torch -line-directive-mode none
+
+[PyExport("Foo")]
+struct Foo
+{
+    TensorView<float> inValues1;
+    TensorView<float> inValues2;
+}
+
+// CUDA: __global__ void __kernel__myKernel(Foo_[[#]] {{[[:alnum:]_]+}}, TensorView {{[[:alnum:]_]+}})
+[AutoPyBindCUDA]
+[CudaKernel]
+void myKernel(Foo foo, TensorView<float> outValues)
+{
+    if (cudaThreadIdx().x > 0)
+        return;
+    outValues.store(
+        cudaThreadIdx().x,
+        sin(foo.inValues1.load(cudaThreadIdx().x)) * cos(foo.inValues2.load(cudaThreadIdx().x)));
+}
+
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: __kernel__myKernel(Foo_[[#]] {{[[:alnum:]_]+}}, TensorView {{[[:alnum:]_]+}});
+// 
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: void myKernel(std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<uint32_t, uint32_t, uint32_t> {{[[:alnum:]_]+}}, std::tuple<torch::Tensor, torch::Tensor> {{[[:alnum:]_]+}}, torch::Tensor {{[[:alnum:]_]+}})
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: std::tuple<std::tuple<const char*, const char*, const char*, const char*>, std::tuple<const char*, const char*>, const char*, const char*> __funcinfo__myKernel()
+//
+// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
+// TORCH-NEXT: std::tuple<std::tuple<const char*, const char*>, std::tuple<const char*, const char*>> __typeinfo__Foo()
+
diff --git a/tests/autodiff/long-loop-branching-addition.slang b/tests/autodiff/long-loop-branching-addition.slang
index beb371bd0..f7e8b43f1 100644
--- a/tests/autodiff/long-loop-branching-addition.slang
+++ b/tests/autodiff/long-loop-branching-addition.slang
@@ -42,7 +42,7 @@ float sin_series(float x, int iterations)
 // This test inparticular checks that can identify induction variables through
 // branching control flow
 
-// CHECK: struct s_bwd_sin_series_Intermediates
+// CHECK: struct s_bwd_prop_sin_series_Intermediates
 // CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/long-loop-chained-addition.slang b/tests/autodiff/long-loop-chained-addition.slang
index 8f75744a9..15bb48aa4 100644
--- a/tests/autodiff/long-loop-chained-addition.slang
+++ b/tests/autodiff/long-loop-chained-addition.slang
@@ -27,7 +27,7 @@ float sin_series(float x, int iterations)
 // This test inparticular checks that can identify induction variables with
 // more than one operation applied to them during the loop
 
-// CHECK: struct s_bwd_sin_series_Intermediates
+// CHECK: struct s_bwd_prop_sin_series_Intermediates
 // CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/long-loop-multiple.slang b/tests/autodiff/long-loop-multiple.slang
index a696beccf..c03226080 100644
--- a/tests/autodiff/long-loop-multiple.slang
+++ b/tests/autodiff/long-loop-multiple.slang
@@ -23,7 +23,7 @@ float sin_series(float x, int iterations)
 // This test differs from ./long-loop.slang in that the loop counter is
 // relative to a multiple of the loop iteration
 
-// CHECK: struct s_bwd_sin_series_Intermediates
+// CHECK: struct s_bwd_prop_sin_series_Intermediates
 // CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/long-loop-noninductive.slang b/tests/autodiff/long-loop-noninductive.slang
index bfd37c4f2..35b5de43b 100644
--- a/tests/autodiff/long-loop-noninductive.slang
+++ b/tests/autodiff/long-loop-noninductive.slang
@@ -25,7 +25,7 @@ float sin_series(float x, int iterations)
 // `i`. This test checks that the induction variable finder doesn't
 // accidentally succeed all the time
 
-// CHECK: struct s_bwd_sin_series_Intermediates
+// CHECK: struct s_bwd_prop_sin_series_Intermediates
 // CHECK: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/long-loop.slang b/tests/autodiff/long-loop.slang
index 69652dbf0..82a9647b1 100644
--- a/tests/autodiff/long-loop.slang
+++ b/tests/autodiff/long-loop.slang
@@ -21,7 +21,7 @@ float sin_series(float x, int iterations)
 
 // Check that the intermediate context of sin_series does not have an array for `i`.
 
-// CHECK: struct s_bwd_sin_series_Intermediates
+// CHECK: struct s_bwd_prop_sin_series_Intermediates
 // CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/long-while-loop.slang b/tests/autodiff/long-while-loop.slang
index 20d802e2a..c329bd757 100644
--- a/tests/autodiff/long-while-loop.slang
+++ b/tests/autodiff/long-while-loop.slang
@@ -25,7 +25,7 @@ float sin_series(float x, int iterations)
 // This differs from ./long-loop.slang in that it uses an equivalent do/while
 // loop, this tests checks that induction variables are still correctly identified.
 
-// CHECK: struct s_bwd_sin_series_Intermediates
+// CHECK: struct s_bwd_prop_sin_series_Intermediates
 // CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/loop-init.slang b/tests/autodiff/loop-init.slang
index 26d837e75..d7afb455c 100644
--- a/tests/autodiff/loop-init.slang
+++ b/tests/autodiff/loop-init.slang
@@ -16,7 +16,7 @@ struct A : IDifferentiable
 // loop state.
 // 
 
-// CHECK: struct s_bwd_B_eval_Intermediates_0
+// CHECK: struct s_bwd_prop_B_eval_Intermediates_0
 // CHECK-NOT: int {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
 
diff --git a/tests/autodiff/reverse-checkpoint-2.slang b/tests/autodiff/reverse-checkpoint-2.slang
index bd787d83a..68ff62176 100644
--- a/tests/autodiff/reverse-checkpoint-2.slang
+++ b/tests/autodiff/reverse-checkpoint-2.slang
@@ -28,7 +28,7 @@ float f(int p, float x)
 
 // Check that there are no calls to primal_g in bwd_f.
 
-// CHECK: void s_bwd_f_{{[0-9]+}}
+// CHECK: void s_bwd_prop_f_{{[0-9]+}}
 // CHECK: {{[_a-zA-Z0-9]+}} = s_bwd_primal_g_{{[0-9]+}}
 // CHECK: return
 
diff --git a/tests/autodiff/reverse-loop-checkpoint-test.slang b/tests/autodiff/reverse-loop-checkpoint-test.slang
index 732360013..fc206e128 100644
--- a/tests/autodiff/reverse-loop-checkpoint-test.slang
+++ b/tests/autodiff/reverse-loop-checkpoint-test.slang
@@ -31,7 +31,7 @@ float infinitesimal(float x)
 
 // Test that computeLoop's intermediates have no float sitting
 // around (must not cache the outvar from 'compute()')
-// CHECK: struct s_bwd_computeLoop_Intermediates
+// CHECK: struct s_bwd_prop_computeLoop_Intermediates
 // CHECK-NEXT: {
 // CHECK-NOT: {{[A-Za-z0-9_]+}} {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
diff --git a/tests/autodiff/reverse-loop-diff-only-2.slang b/tests/autodiff/reverse-loop-diff-only-2.slang
index aad405b48..2cc33ecca 100644
--- a/tests/autodiff/reverse-loop-diff-only-2.slang
+++ b/tests/autodiff/reverse-loop-diff-only-2.slang
@@ -32,7 +32,7 @@ float infinitesimal(float x)
 
 // Test that computeLoop's intermediates have no float sitting
 // around (must not cache the outvar from 'compute()')
-// CHECK: struct s_bwd_computeLoop_Intermediates
+// CHECK: struct s_bwd_prop_computeLoop_Intermediates
 // CHECK-NEXT: {
 // CHECK-NOT: {{[A-Za-z0-9_]+}} {{[A-Za-z0-9_]+}}[{{.*}}]
 // CHECK: }
diff --git a/tests/compute/generic-default-arg.slang b/tests/compute/generic-default-arg.slang
new file mode 100644
index 000000000..8762f5e8a
--- /dev/null
+++ b/tests/compute/generic-default-arg.slang
@@ -0,0 +1,53 @@
+//TEST(compute):COMPARE_COMPUTE: -shaderobj
+//TEST(compute):COMPARE_COMPUTE:-cpu -shaderobj
+
+// Check that user code can declare and use a generic
+// `struct` type.
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+interface ITest
+{
+	int doThing(int x);
+};
+
+struct Impl1 : ITest
+{
+	int doThing(int x)
+	{
+		return x * 2;
+	}
+};
+
+struct Impl2 : ITest
+{
+    int doThing(int x)
+    {
+        return x * 3;
+    }
+};
+
+__generic<T : ITest = Impl1>
+struct GenStruct
+{
+	T obj;
+};
+
+int test(GenStruct gs, int val)
+{
+	return gs.obj.doThing(val);
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+	int tid = dispatchThreadID.x;
+
+	int outVal = 0;
+
+	GenStruct<Impl1> gs;
+	outVal += test(gs, tid);
+
+	outputBuffer[tid] = outVal;
+}
\ No newline at end of file
diff --git a/tests/compute/generic-default-arg.slang.expected.txt b/tests/compute/generic-default-arg.slang.expected.txt
new file mode 100644
index 000000000..e1e8ccec4
--- /dev/null
+++ b/tests/compute/generic-default-arg.slang.expected.txt
@@ -0,0 +1,4 @@
+0
+2
+4
+6
diff --git a/tests/diagnostics/autodiff-data-flow.slang.expected b/tests/diagnostics/autodiff-data-flow.slang.expected
index 301f84985..6840bfd3c 100644
--- a/tests/diagnostics/autodiff-data-flow.slang.expected
+++ b/tests/diagnostics/autodiff-data-flow.slang.expected
@@ -3,9 +3,6 @@ standard error = {
 tests/diagnostics/autodiff-data-flow.slang(15): error 41020: derivative cannot be propagated through call to non-forward-differentiable function `nonDiff`, use 'no_diff' to clarify intention.
         val = nonDiff(x * 2.0f);
                      ^
-tests/diagnostics/autodiff-data-flow.slang(22): error 41021: a differentiable function must have at least one differentiable output.
-void g(float x)
-     ^
 tests/diagnostics/autodiff-data-flow.slang(28): error 30510: loops inside a differentiable function need to provide either '[MaxIters(n)]' or '[ForceUnroll]' attribute.
     for (int i = 0; i < 5; i++) // Not ok, we can't infer the loop iterations because the body modifies induction var.
     ^~~
diff --git a/tests/diagnostics/generic-incorrect-default-arg.slang b/tests/diagnostics/generic-incorrect-default-arg.slang
new file mode 100644
index 000000000..14192293c
--- /dev/null
+++ b/tests/diagnostics/generic-incorrect-default-arg.slang
@@ -0,0 +1,53 @@
+//TEST:SIMPLE(filecheck=CHECK): -target hlsl -profile cs_5_0 -entry computeMain -line-directive-mode none
+
+// Check that user code can declare and use a generic
+// `struct` type.
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+interface ITest
+{
+	int doThing(int x);
+};
+
+struct Impl1 : ITest
+{
+	int doThing(int x)
+	{
+		return x * 2;
+	}
+};
+
+struct Impl2
+{
+    int doSomethingElse(int x)
+    {
+        return x * 3;
+    }
+};
+
+__generic<T : ITest = Impl2>
+// CHECK: tests/diagnostics/generic-incorrect-default-arg.slang([[@LINE-1]]): error 38029: type argument 'Impl2' does not conform to the required interface 'ITest'
+struct GenStruct
+{
+	T obj;
+};
+
+int test(GenStruct gs, int val)
+{
+	return gs.obj.doThing(val);
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+	int tid = dispatchThreadID.x;
+
+	int outVal = 0;
+
+	GenStruct<Impl1> gs;
+	outVal += test(gs, tid);
+
+	outputBuffer[tid] = outVal;
+}
\ No newline at end of file