Skip to content

Commit

Permalink
aded device_log!
Browse files Browse the repository at this point in the history
  • Loading branch information
shiinamiyuki committed Nov 4, 2023
1 parent bad842e commit 255fa72
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 42 deletions.
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use std::env::current_exe;

use luisa::prelude::*;
use luisa::printer::*;

use luisa_compute as luisa;

Expand All @@ -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);
});
}
68 changes: 35 additions & 33 deletions luisa_compute/src/lang/print.rs
Original file line number Diff line number Diff line change
@@ -1,55 +1,57 @@
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(),
args: Vec::new(),
}
}
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();
}};
}
7 changes: 5 additions & 2 deletions luisa_compute/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 255fa72

Please sign in to comment.