From ed7d5a6f7aac0631bd10876e7939285f53b5c511 Mon Sep 17 00:00:00 2001
From: ReversedCausality <rg@youxplode.com>
Date: Tue, 3 Oct 2023 21:31:10 +0100
Subject: [PATCH 1/3] Made `build_kernel` public.

---
 luisa_compute/src/runtime/kernel.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/luisa_compute/src/runtime/kernel.rs b/luisa_compute/src/runtime/kernel.rs
index a39d9f3..2b6088e 100644
--- a/luisa_compute/src/runtime/kernel.rs
+++ b/luisa_compute/src/runtime/kernel.rs
@@ -381,7 +381,7 @@ impl KernelBuilder {
             }
         })
     }
-    fn build_kernel<S: KernelSignature>(
+    pub fn build_kernel<S: KernelSignature>(
         &mut self,
         body: impl FnOnce(&mut Self),
     ) -> crate::runtime::KernelDef<S> {

From 1d9baf0c9daa8ade1d697782c30b9710ee2fa6e0 Mon Sep 17 00:00:00 2001
From: ReversedCausality <rg@youxplode.com>
Date: Fri, 6 Oct 2023 19:10:41 +0100
Subject: [PATCH 2/3] Updated AsKernelArg to provide output.

---
 luisa_compute/src/runtime.rs | 65 ++++++++++++++++++------------------
 1 file changed, 33 insertions(+), 32 deletions(-)

diff --git a/luisa_compute/src/runtime.rs b/luisa_compute/src/runtime.rs
index c8f1491..603522c 100644
--- a/luisa_compute/src/runtime.rs
+++ b/luisa_compute/src/runtime.rs
@@ -1375,45 +1375,46 @@ impl<T: KernelSignature> Kernel<T> {
     }
 }
 
-pub trait AsKernelArg<T: KernelArg>: KernelArg {}
-
-impl<T: Value> AsKernelArg<T> for T {}
-
-impl<T: Value> AsKernelArg<Buffer<T>> for Buffer<T> {}
-
-impl<'a, T: Value> AsKernelArg<Buffer<T>> for BufferView<'a, T> {}
-
-impl<'a, T: Value> AsKernelArg<BufferView<'a, T>> for BufferView<'a, T> {}
-
-impl<'a, T: Value> AsKernelArg<BufferView<'a, T>> for Buffer<T> {}
-
-// impl AsKernelArg<ByteBuffer> for ByteBuffer {}
-
-// impl<'a> AsKernelArg<ByteBuffer> for ByteBufferView<'a> {}
-
-// impl<'a> AsKernelArg<ByteBufferView<'a>> for ByteBufferView<'a> {}
-
-// impl<'a> AsKernelArg<ByteBufferView<'a>> for ByteBuffer {}
-
-impl<'a, T: IoTexel> AsKernelArg<Tex2d<T>> for Tex2dView<'a, T> {}
+// A trait signifying that this argument can be used in place of an argument of type `Self::T`.
+pub trait AsKernelArg: KernelArg {
+    type Output: KernelArg + 'static;
+}
 
-impl<'a, T: IoTexel> AsKernelArg<Tex3d<T>> for Tex3dView<'a, T> {}
+impl<T: Value> AsKernelArg for T {
+    type Output = T;
+}
 
-impl<'a, T: IoTexel> AsKernelArg<Tex3dView<'a, T>> for Tex2dView<'a, T> {}
+impl<T: Value> AsKernelArg for Buffer<T> {
+    type Output = T;
+}
 
-impl<'a, T: IoTexel> AsKernelArg<Tex3dView<'a, T>> for Tex3dView<'a, T> {}
+impl<'a, T: Value> AsKernelArg for BufferView<'a, T> {
+    type Output = Buffer<T>;
+}
 
-impl<'a, T: IoTexel> AsKernelArg<Tex3dView<'a, T>> for Tex2d<T> {}
+impl<'a, T: IoTexel> AsKernelArg for Tex2dView<'a, T> {
+    type Output = Tex2d<T>;
+}
 
-impl<'a, T: IoTexel> AsKernelArg<Tex3dView<'a, T>> for Tex3d<T> {}
+impl<'a, T: IoTexel> AsKernelArg for Tex3dView<'a, T> {
+    type Output = Tex3d<T>;
+}
 
-impl<T: IoTexel> AsKernelArg<Tex2d<T>> for Tex2d<T> {}
+impl<T: IoTexel> AsKernelArg for Tex2d<T> {
+    type Output = Tex2d<T>;
+}
 
-impl<T: IoTexel> AsKernelArg<Tex3d<T>> for Tex3d<T> {}
+impl<T: IoTexel> AsKernelArg for Tex3d<T> {
+    type Output = Tex3d<T>;
+}
 
-impl AsKernelArg<BindlessArray> for BindlessArray {}
+impl AsKernelArg for BindlessArray {
+    type Output = BindlessArray;
+}
 
-impl AsKernelArg<Accel> for Accel {}
+impl AsKernelArg for Accel {
+    type Output = Accel;
+}
 
 macro_rules! impl_call_for_callable {
     ( $($Ts:ident)*) => {
@@ -1463,7 +1464,7 @@ macro_rules! impl_dispatch_for_kernel {
         impl <$($Ts: KernelArg+'static),*> Kernel<fn($($Ts,)*)> {
             #[allow(non_snake_case)]
             #[allow(unused_mut)]
-            pub fn dispatch(&self, dispatch_size: [u32; 3], $($Ts:&impl AsKernelArg<$Ts>),*)  {
+            pub fn dispatch(&self, dispatch_size: [u32; 3], $($Ts:&impl AsKernelArg<Output = $Ts>),*)  {
                 let mut encoder = KernelArgEncoder::new();
                 $($Ts.encode(&mut encoder);)*
                 self.inner.dispatch(encoder, dispatch_size)
@@ -1472,7 +1473,7 @@ macro_rules! impl_dispatch_for_kernel {
             #[allow(unused_mut)]
             pub fn dispatch_async(
                 &self,
-                dispatch_size: [u32; 3], $($Ts:&impl AsKernelArg<$Ts>),*
+                dispatch_size: [u32; 3], $($Ts:&impl AsKernelArg<Output = $Ts>),*
             ) -> Command<'static, 'static> {
                 let mut encoder = KernelArgEncoder::new();
                 $($Ts.encode(&mut encoder);)*

From 9e8097188e5a209065ebdaba10ed85006adb95e9 Mon Sep 17 00:00:00 2001
From: ReversedCausality <rg@youxplode.com>
Date: Fri, 6 Oct 2023 19:19:44 +0100
Subject: [PATCH 3/3] Fix macro.

---
 luisa_compute_derive_impl/src/lib.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/luisa_compute_derive_impl/src/lib.rs b/luisa_compute_derive_impl/src/lib.rs
index 3791fd2..c40e21c 100644
--- a/luisa_compute_derive_impl/src/lib.rs
+++ b/luisa_compute_derive_impl/src/lib.rs
@@ -141,7 +141,8 @@ impl Compiler {
                     #(self.#field_names.encode(encoder);)*
                 }
             }
-            impl #impl_generics #runtime_path::AsKernelArg<#name #ty_generics> for #name #ty_generics #where_clause {
+            impl #impl_generics #runtime_path::AsKernelArg for #name #ty_generics #where_clause {
+                type Output = #name #ty_generics;
             }
         )
     }