From 96049efe42340b7c4c41a092b65f1966cebbf1cc Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Mon, 4 Dec 2023 16:03:09 -0500
Subject: [PATCH] added address() method for interop with external code

---
 luisa_compute/src/lang/types.rs | 12 ++++++++++++
 luisa_compute/src/resource.rs   | 24 ++++++++++++++++++++++++
 luisa_compute_sys/LuisaCompute  |  2 +-
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/luisa_compute/src/lang/types.rs b/luisa_compute/src/lang/types.rs
index 15de9c9..633fddb 100644
--- a/luisa_compute/src/lang/types.rs
+++ b/luisa_compute/src/lang/types.rs
@@ -317,6 +317,12 @@ impl<T: Value> Expr<T> {
     pub fn _type_tag(&self) -> TypeTag<T> {
         TypeTag(PhantomData)
     }
+    /// Address of this `Expr`, only available in some backends.
+    pub fn address(&self) -> Expr<u64> {
+        let self_node = self.node().get();
+        let node = __current_scope(|b| b.call(Func::AddressOf, &[self_node], u64::type_()));
+        Expr::<u64>::from_node(node.into())
+    }
 }
 
 impl<T: Value> Var<T> {
@@ -344,6 +350,12 @@ impl<T: Value> Var<T> {
     pub fn store(&self, value: impl AsExpr<Value = T>) {
         crate::lang::_store(self, &value.as_expr());
     }
+    /// Address of this `Var`, only available in some backends.
+    pub fn address(&self) -> Expr<u64> {
+        let self_node = self.node().get();
+        let node = __current_scope(|b| b.call(Func::AddressOf, &[self_node], u64::type_()));
+        Expr::<u64>::from_node(node.into())
+    }
 }
 impl<T: Value> AtomicRef<T> {
     pub fn _ref<'a>(self) -> &'a Self {
diff --git a/luisa_compute/src/resource.rs b/luisa_compute/src/resource.rs
index c23960a..5e71d1f 100644
--- a/luisa_compute/src/resource.rs
+++ b/luisa_compute/src/resource.rs
@@ -1398,6 +1398,14 @@ pub struct BufferVar<T: Value> {
     pub(crate) handle: Option<Arc<BufferHandle>>,
     pub(crate) node: SafeNodeRef,
 }
+impl<T: Value> BufferVar<T> {
+    pub fn device_address(&self) -> Expr<u64> {
+        let buf = self.node().get();
+        Expr::from_node(
+            __current_scope(|b| b.call(Func::BufferAddress, &[buf], u64::type_())).into(),
+        )
+    }
+}
 impl<T: Value> ToNode for BufferVar<T> {
     fn node(&self) -> SafeNodeRef {
         self.node
@@ -1423,6 +1431,22 @@ impl<T: Value> ToNode for BindlessBufferVar<T> {
         self.array
     }
 }
+impl<T: Value> BindlessBufferVar<T> {
+    pub fn device_address(&self) -> Expr<u64> {
+        let array = self.array.get();
+        let buffer_index = self.buffer_index.node().get();
+        Expr::from_node(
+            __current_scope(|b| {
+                b.call(
+                    Func::BindlessBufferAddress,
+                    &[array, buffer_index],
+                    u64::type_(),
+                )
+            })
+            .into(),
+        )
+    }
+}
 
 impl<T: Value> IndexRead for BindlessBufferVar<T> {
     type Element = T;
diff --git a/luisa_compute_sys/LuisaCompute b/luisa_compute_sys/LuisaCompute
index 4b2f0c2..0e55e64 160000
--- a/luisa_compute_sys/LuisaCompute
+++ b/luisa_compute_sys/LuisaCompute
@@ -1 +1 @@
-Subproject commit 4b2f0c2492f33c78a37008de96d555aaaf9e33a2
+Subproject commit 0e55e64645e168ae9e0cbd6434983d5c0e026ce2