added address() method for interop with external code

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 {

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;