From 255fa72dac829b31473af9781cb7e0594b874102 Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Sat, 4 Nov 2023 02:25:30 -0400
Subject: [PATCH] aded device_log!

---
 .../examples/{printer.rs => device_print.rs}  |  8 +--
 luisa_compute/src/lang/print.rs               | 68 ++++++++++---------
 luisa_compute/src/lib.rs                      |  7 +-
 luisa_compute_sys/LuisaCompute                |  2 +-
 4 files changed, 43 insertions(+), 42 deletions(-)
 rename luisa_compute/examples/{printer.rs => device_print.rs} (75%)

diff --git a/luisa_compute/examples/printer.rs b/luisa_compute/examples/device_print.rs
similarity index 75%
rename from luisa_compute/examples/printer.rs
rename to luisa_compute/examples/device_print.rs
index ea070dc..ff3afa3 100644
--- a/luisa_compute/examples/printer.rs
+++ b/luisa_compute/examples/device_print.rs
@@ -1,7 +1,6 @@
 use std::env::current_exe;
 
 use luisa::prelude::*;
-use luisa::printer::*;
 
 use luisa_compute as luisa;
 
@@ -20,21 +19,18 @@ fn main() {
     } else {
         "cpu"
     });
-    let printer = Printer::new(&device, 65536);
     let kernel = Kernel::<fn()>::new(
         &device,
         &track!(|| {
             let id = dispatch_id().xy();
             if id.x == id.y {
-                lc_info!(printer, "id = {:?}", id);
+                device_log!("id = {}", id);
             } else {
-                lc_info!(printer, "not equal!, id = [{} {}]", id.x, id.y);
+                device_log!("not equal!, id = [{} {}]", id.x, id.y);
             }
         }),
     );
     device.default_stream().with_scope(|s| {
-        s.reset_printer(&printer);
         s.submit([kernel.dispatch_async([4, 4, 1])]);
-        s.print(&printer);
     });
 }
diff --git a/luisa_compute/src/lang/print.rs b/luisa_compute/src/lang/print.rs
index 8f223fa..d749ccb 100644
--- a/luisa_compute/src/lang/print.rs
+++ b/luisa_compute/src/lang/print.rs
@@ -1,12 +1,12 @@
-use crate::internal_prelude::*;
-
-use super::types::core::Primitive;
+use std::ffi::CString;
 
-pub struct DebugFormatter {
+use crate::internal_prelude::*;
+/// TODO: support custom print format
+pub struct DevicePrintFormatter {
     pub(crate) fmt: String,
     pub(crate) args: Vec<NodeRef>,
 }
-impl DebugFormatter {
+impl DevicePrintFormatter {
     pub fn new() -> Self {
         Self {
             fmt: String::new(),
@@ -14,42 +14,44 @@ impl DebugFormatter {
         }
     }
     pub fn push_str(&mut self, s: &str) {
-        assert!(
-            s.find("{}").is_none(),
-            "DebugFormatter::push_str cannot contain {{}}"
-        );
         self.fmt.push_str(s);
     }
-    pub fn push_arg(&mut self, node: NodeRef) {
-        assert!(
-            node.type_().is_primitive(),
-            "DebugFormatter::push_arg must be primitive"
-        );
-        self.fmt.push_str("{}");
-        self.args.push(node);
+    pub fn push_arg(&mut self, node: SafeNodeRef) {
+        self.args.push(node.get());
     }
-}
-pub trait DebugPrintValue: Value {
-    fn fmt_args(e: Expr<Self>, fmt: &mut DebugFormatter);
-}
-
-impl<T: DebugPrintValue + Primitive> DebugPrintValue for T {
-    fn fmt_args(e: Expr<Self>, fmt: &mut DebugFormatter) {
-        fmt.push_arg(e.node().get());
+    pub fn print(self) {
+        let Self { fmt, args } = self;
+        __current_scope(|b| {
+            b.print(
+                CBoxedSlice::new(CString::new(fmt).unwrap().into_bytes_with_nul()),
+                &args,
+            );
+        })
     }
 }
-
-pub trait DebugPrint {
-    fn fmt_args(&self, fmt: &mut DebugFormatter);
+pub trait DevicePrint {
+    fn fmt(&self, fmt: &mut DevicePrintFormatter);
 }
 
-impl<T: DebugPrintValue> DebugPrint for Expr<T> {
-    fn fmt_args(&self, fmt: &mut DebugFormatter) {
-        T::fmt_args(*self, fmt);
+impl<T: Value> DevicePrint for Expr<T> {
+    fn fmt(&self, fmt: &mut DevicePrintFormatter) {
+        fmt.push_arg(self.node());
     }
 }
-impl<T: DebugPrintValue> DebugPrint for Var<T> {
-    fn fmt_args(&self, fmt: &mut DebugFormatter) {
-        T::fmt_args(self.load(), fmt);
+impl<T: Value> DevicePrint for Var<T> {
+    fn fmt(&self, fmt: &mut DevicePrintFormatter) {
+        fmt.push_arg(self.node());
     }
 }
+
+#[macro_export]
+macro_rules! device_log {
+    ($fmt:literal, $($arg:expr),*) => {{
+        let mut fmt = $crate::lang::print::DevicePrintFormatter::new();
+        fmt.push_str($fmt);
+        $(
+            $crate::lang::print::DevicePrint::fmt(&$arg, &mut fmt);
+        )*
+        fmt.print();
+    }};
+}
diff --git a/luisa_compute/src/lib.rs b/luisa_compute/src/lib.rs
index 21e70d9..5ffc74f 100644
--- a/luisa_compute/src/lib.rs
+++ b/luisa_compute/src/lib.rs
@@ -44,7 +44,8 @@ pub mod prelude {
         Stream, Swapchain,
     };
     pub use crate::{
-        cpu_dbg, if_, lc_assert, lc_comment_lineno, lc_unreachable, loop_, while_, Context,
+        cpu_dbg, device_log, if_, lc_assert, lc_comment_lineno, lc_unreachable, loop_, while_,
+        Context,
     };
 
     pub use luisa_compute_derive::*;
@@ -187,7 +188,9 @@ impl ResourceTracker {
     pub fn upgrade(&self) -> Self {
         let mut strong_refs = vec![];
         for r in self.weak_refs.iter() {
-            strong_refs.push(r.upgrade().unwrap_or_else(|| panic!("Bad weak ref. Kernel captured resources might be dropped.")));
+            strong_refs.push(r.upgrade().unwrap_or_else(|| {
+                panic!("Bad weak ref. Kernel captured resources might be dropped.")
+            }));
         }
         strong_refs.extend(self.strong_refs.iter().cloned());
         Self {
diff --git a/luisa_compute_sys/LuisaCompute b/luisa_compute_sys/LuisaCompute
index f506513..1dce6c3 160000
--- a/luisa_compute_sys/LuisaCompute
+++ b/luisa_compute_sys/LuisaCompute
@@ -1 +1 @@
-Subproject commit f506513e3728518e8c3f1fab93cfa196b7eed265
+Subproject commit 1dce6c31d72a6416941c20c1516e4a8b1df20a78