dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix cube docs

This commit is contained in:
nathaniel 2024-07-12 09:25:45 -04:00
parent 19f5ad7be5
commit 0a33aa363d
3 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
crates/burn-cube-macros/src/codegen_function/launch.rs
View File

@ -463,8 +463,7 @@ pub fn codegen_launch(sig: &syn::Signature) -> TokenStream {
let kernel = codegen.gen_kernel_struct(); let kernel = codegen.gen_kernel_struct();
let compile = codegen.gen_compile_impl(&ident_expand); let compile = codegen.gen_compile_impl(&ident_expand);
let (inputs, output) = (codegen.fn_inputs, codegen.fn_output); let (inputs, output) = (codegen.fn_inputs, codegen.fn_output);
let doc = let doc = format!("Launch the kernel [{ident}()] on the given runtime.");
format!("Launch the kernel [{ident}] with the provided argument on the given runtime.");
quote::quote! { quote::quote! {
#kernel #kernel
@ -472,7 +471,6 @@ pub fn codegen_launch(sig: &syn::Signature) -> TokenStream {
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
#[doc = #doc] #[doc = #doc]
/// Launch the kernel.
pub fn launch #generics ( pub fn launch #generics (
client: ComputeClient<R::Server, R::Channel>, client: ComputeClient<R::Server, R::Channel>,
cube_count: CubeCount<R::Server>, cube_count: CubeCount<R::Server>,

2
crates/burn-cube-macros/src/codegen_trait/mod.rs
View File

@ -103,7 +103,7 @@ fn register_expand(
quote::quote! ( quote::quote! (
#expand { #expand {
#[cube] #[cube]
pub #func #func
#func_expand #func_expand
} }
) )

8
crates/burn-cube-macros/src/lib.rs
View File

@ -143,7 +143,11 @@ fn codegen_cube(
return Err(code); return Err(code);
} }
let launch_doc = if launch { "and launch function " } else { "" }; let launch_doc = if launch {
"and launch functions "
} else {
"function "
};
let launch = if launch { let launch = if launch {
codegen_launch(&func.sig) codegen_launch(&func.sig)
@ -153,7 +157,7 @@ fn codegen_cube(
let mod_name = &func.sig.ident; let mod_name = &func.sig.ident;
let vis = &func.vis; let vis = &func.vis;
let doc = format!("Module containing the expand method {launch_doc}of {mod_name}."); let doc = format!("Module containing the expand {launch_doc}of {mod_name}.");
Ok(quote::quote! { Ok(quote::quote! {
#[allow(dead_code)] #[allow(dead_code)]