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

13 files changed, 114 insertions, 10 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: }