Added lint checking support to monomorphised kernel impls

juntyr · juntyr · commit f37ef1af305c · 2023-08-10T21:46:48.000Z
diff --git a/rust-cuda-derive/src/kernel/link/config.rs b/rust-cuda-derive/src/kernel/link/config.rs
@@ -1,4 +1,6 @@
-use std::path::PathBuf;
+use std::{collections::HashMap, path::PathBuf};
+
+use super::super::lints::{parse_ptx_lint_level, LintLevel, PtxLint};
 
 #[allow(clippy::module_name_repetitions)]
 pub(super) struct LinkKernelConfig {
@@ -8,6 +10,7 @@ pub(super) struct LinkKernelConfig {
     pub(super) crate_name: String,
     pub(super) crate_path: PathBuf,
     pub(super) specialisation: String,
+    pub(super) ptx_lint_levels: HashMap<PtxLint, LintLevel>,
 }
 
 impl syn::parse::Parse for LinkKernelConfig {
@@ -37,13 +40,27 @@ impl syn::parse::Parse for LinkKernelConfig {
             String::new()
         };
 
+        let attrs = syn::punctuated::Punctuated::<
+            syn::MetaList,
+            syn::token::Comma,
+        >::parse_separated_nonempty(input)?;
+
+        let mut ptx_lint_levels = HashMap::new();
+
+        for syn::MetaList { path, nested, .. } in attrs {
+            parse_ptx_lint_level(&path, &nested, &mut ptx_lint_levels);
+        }
+
+        proc_macro_error::abort_if_dirty();
+
         Ok(Self {
             kernel,
             kernel_hash,
             args,
             crate_name: name.value(),
             crate_path: PathBuf::from(path.value()),
             specialisation,
+            ptx_lint_levels,
         })
     }
 }
diff --git a/rust-cuda-derive/src/kernel/link/mod.rs b/rust-cuda-derive/src/kernel/link/mod.rs
@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     env,
     ffi::CString,
     fmt::Write as FmtWrite,
@@ -17,7 +18,10 @@ use ptx_builder::{
     error::{BuildErrorKind, Error, Result},
 };
 
-use super::utils::skip_kernel_compilation;
+use super::{
+    lints::{LintLevel, PtxLint},
+    utils::skip_kernel_compilation,
+};
 
 mod config;
 mod error;
@@ -68,12 +72,14 @@ pub fn link_kernel(tokens: TokenStream) -> TokenStream {
         crate_name,
         crate_path,
         specialisation,
+        ptx_lint_levels,
     } = match syn::parse_macro_input::parse(tokens) {
         Ok(config) => config,
         Err(err) => {
             abort_call_site!(
-                "link_kernel!(KERNEL ARGS NAME PATH SPECIALISATION) expects KERNEL and ARGS \
-                 identifiers, NAME and PATH string literals, and SPECIALISATION tokens: {:?}",
+                "link_kernel!(KERNEL ARGS NAME PATH SPECIALISATION LINTS,*) expects KERNEL and \
+                 ARGS identifiers, NAME and PATH string literals, SPECIALISATION and LINTS \
+                 tokens: {:?}",
                 err
             )
         },
@@ -201,7 +207,7 @@ pub fn link_kernel(tokens: TokenStream) -> TokenStream {
     }
 
     let (result, error_log, info_log, version, drop) =
-        check_kernel_ptx(&kernel_ptx, &specialisation, &kernel_hash);
+        check_kernel_ptx(&kernel_ptx, &specialisation, &kernel_hash, &ptx_lint_levels);
 
     let ptx_compiler = match &version {
         Ok((major, minor)) => format!("PTX compiler v{major}.{minor}"),
@@ -272,10 +278,12 @@ pub fn link_kernel(tokens: TokenStream) -> TokenStream {
 }
 
 #[allow(clippy::type_complexity)]
+#[allow(clippy::too_many_lines)]
 fn check_kernel_ptx(
     kernel_ptx: &str,
     specialisation: &str,
     kernel_hash: &proc_macro2::Ident,
+    ptx_lint_levels: &HashMap<PtxLint, LintLevel>,
 ) -> (
     Result<(), NvptxError>,
     Result<Option<String>, NvptxError>,
@@ -297,7 +305,7 @@ fn check_kernel_ptx(
         compiler
     };
 
-    let result = {
+    let result = (|| {
         let kernel_name = if specialisation.is_empty() {
             format!("{kernel_hash}_kernel")
         } else {
@@ -306,15 +314,79 @@ fn check_kernel_ptx(
                 seahash::hash(specialisation.as_bytes())
             )
         };
-
-        let options = vec![
+        let mut options = vec![
             CString::new("--entry").unwrap(),
             CString::new(kernel_name).unwrap(),
-            CString::new("--verbose").unwrap(),
-            CString::new("--warn-on-double-precision-use").unwrap(),
-            CString::new("--warn-on-local-memory-usage").unwrap(),
-            CString::new("--warn-on-spills").unwrap(),
         ];
+
+        if ptx_lint_levels
+            .values()
+            .any(|level| *level > LintLevel::Warn)
+        {
+            let mut options = options.clone();
+
+            if ptx_lint_levels
+                .get(&PtxLint::Verbose)
+                .map_or(false, |level| *level > LintLevel::Warn)
+            {
+                options.push(CString::new("--verbose").unwrap());
+            }
+            if ptx_lint_levels
+                .get(&PtxLint::DoublePrecisionUse)
+                .map_or(false, |level| *level > LintLevel::Warn)
+            {
+                options.push(CString::new("--warn-on-double-precision-use").unwrap());
+            }
+            if ptx_lint_levels
+                .get(&PtxLint::LocalMemoryUsage)
+                .map_or(false, |level| *level > LintLevel::Warn)
+            {
+                options.push(CString::new("--warn-on-local-memory-usage").unwrap());
+            }
+            if ptx_lint_levels
+                .get(&PtxLint::RegisterSpills)
+                .map_or(false, |level| *level > LintLevel::Warn)
+            {
+                options.push(CString::new("--warn-on-spills").unwrap());
+            }
+            options.push(CString::new("--warning-as-error").unwrap());
+
+            let options_ptrs = options.iter().map(|o| o.as_ptr()).collect::<Vec<_>>();
+
+            NvptxError::try_err_from(unsafe {
+                ptx_compiler_sys::nvPTXCompilerCompile(
+                    compiler,
+                    options_ptrs.len() as c_int,
+                    options_ptrs.as_ptr().cast(),
+                )
+            })?;
+        };
+
+        if ptx_lint_levels
+            .get(&PtxLint::Verbose)
+            .map_or(false, |level| *level > LintLevel::Allow)
+        {
+            options.push(CString::new("--verbose").unwrap());
+        }
+        if ptx_lint_levels
+            .get(&PtxLint::DoublePrecisionUse)
+            .map_or(false, |level| *level > LintLevel::Allow)
+        {
+            options.push(CString::new("--warn-on-double-precision-use").unwrap());
+        }
+        if ptx_lint_levels
+            .get(&PtxLint::LocalMemoryUsage)
+            .map_or(false, |level| *level > LintLevel::Allow)
+        {
+            options.push(CString::new("--warn-on-local-memory-usage").unwrap());
+        }
+        if ptx_lint_levels
+            .get(&PtxLint::RegisterSpills)
+            .map_or(false, |level| *level > LintLevel::Allow)
+        {
+            options.push(CString::new("--warn-on-spills").unwrap());
+        }
+
         let options_ptrs = options.iter().map(|o| o.as_ptr()).collect::<Vec<_>>();
 
         NvptxError::try_err_from(unsafe {
@@ -324,7 +396,7 @@ fn check_kernel_ptx(
                 options_ptrs.as_ptr().cast(),
             )
         })
-    };
+    })();
 
     let error_log = (|| {
         let mut error_log_size = 0;
diff --git a/rust-cuda-derive/src/kernel/lints.rs b/rust-cuda-derive/src/kernel/lints.rs
@@ -0,0 +1,154 @@
+use std::{collections::HashMap, fmt};
+
+use syn::spanned::Spanned;
+
+pub fn parse_ptx_lint_level(
+    path: &syn::Path,
+    nested: &syn::punctuated::Punctuated<syn::NestedMeta, syn::token::Comma>,
+    ptx_lint_levels: &mut HashMap<PtxLint, LintLevel>,
+) {
+    let level = match path.get_ident() {
+        Some(ident) if ident == "allow" => LintLevel::Allow,
+        Some(ident) if ident == "warn" => LintLevel::Warn,
+        Some(ident) if ident == "deny" => LintLevel::Deny,
+        Some(ident) if ident == "forbid" => LintLevel::Forbid,
+        _ => {
+            emit_error!(
+                path.span(),
+                "[rust-cuda]: Invalid lint #[kernel(<level>(<lint>))] attribute: unknown lint \
+                 level, must be one of `allow`, `warn`, `deny`, `forbid`.",
+            );
+
+            return;
+        },
+    };
+
+    for meta in nested {
+        let syn::NestedMeta::Meta(syn::Meta::Path(path)) = meta else {
+            emit_error!(
+                meta.span(),
+                "[rust-cuda]: Invalid #[kernel({}(<lint>))] attribute.",
+                level,
+            );
+            continue;
+        };
+
+        if path.leading_colon.is_some()
+            || path.segments.empty_or_trailing()
+            || path.segments.len() != 2
+        {
+            emit_error!(
+                meta.span(),
+                "[rust-cuda]: Invalid #[kernel({}(<lint>))] attribute: <lint> must be of the form \
+                 `ptx::lint`.",
+                level,
+            );
+            continue;
+        }
+
+        let Some(syn::PathSegment { ident: namespace, arguments: syn::PathArguments::None }) = path.segments.first() else {
+            emit_error!(
+                meta.span(),
+                "[rust-cuda]: Invalid #[kernel({}(<lint>))] attribute: <lint> must be of the form `ptx::lint`.",
+                level,
+            );
+            continue;
+        };
+
+        if namespace != "ptx" {
+            emit_error!(
+                meta.span(),
+                "[rust-cuda]: Invalid #[kernel({}(<lint>))] attribute: <lint> must be of the form \
+                 `ptx::lint`.",
+                level,
+            );
+            continue;
+        }
+
+        let Some(syn::PathSegment { ident: lint, arguments: syn::PathArguments::None }) = path.segments.last() else {
+            emit_error!(
+                meta.span(),
+                "[rust-cuda]: Invalid #[kernel({}(<lint>))] attribute: <lint> must be of the form `ptx::lint`.",
+                level,
+            );
+            continue;
+        };
+
+        let lint = match lint {
+            l if l == "verbose" => PtxLint::Verbose,
+            l if l == "double_precision_use" => PtxLint::DoublePrecisionUse,
+            l if l == "local_memory_usage" => PtxLint::LocalMemoryUsage,
+            l if l == "register_spills" => PtxLint::RegisterSpills,
+            _ => {
+                emit_error!(
+                    meta.span(),
+                    "[rust-cuda]: Unknown PTX kernel lint `ptx::{}`.",
+                    lint,
+                );
+                continue;
+            },
+        };
+
+        match ptx_lint_levels.get(&lint) {
+            None => (),
+            Some(LintLevel::Forbid) if level < LintLevel::Forbid => {
+                emit_error!(
+                    meta.span(),
+                    "[rust-cuda]: {}(ptx::{}) incompatible with previous forbid.",
+                    level,
+                    lint,
+                );
+                continue;
+            },
+            Some(previous) => {
+                emit_warning!(
+                    meta.span(),
+                    "[rust-cuda]: {}(ptx::{}) overwrites previous {}.",
+                    level,
+                    lint,
+                    previous,
+                );
+            },
+        }
+
+        ptx_lint_levels.insert(lint, level);
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
+pub enum LintLevel {
+    Allow,
+    Warn,
+    Deny,
+    Forbid,
+}
+
+impl fmt::Display for LintLevel {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::Allow => fmt.write_str("allow"),
+            Self::Warn => fmt.write_str("warn"),
+            Self::Deny => fmt.write_str("deny"),
+            Self::Forbid => fmt.write_str("forbid"),
+        }
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
+pub enum PtxLint {
+    Verbose,
+    DoublePrecisionUse,
+    LocalMemoryUsage,
+    RegisterSpills,
+}
+
+impl fmt::Display for PtxLint {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::Verbose => fmt.write_str("verbose"),
+            Self::DoublePrecisionUse => fmt.write_str("double_precision_use"),
+            Self::LocalMemoryUsage => fmt.write_str("local_memory_usage"),
+            Self::RegisterSpills => fmt.write_str("register_spills"),
+        }
+    }
+}
diff --git a/rust-cuda-derive/src/kernel/mod.rs b/rust-cuda-derive/src/kernel/mod.rs
@@ -2,4 +2,5 @@ pub mod link;
 pub mod specialise;
 pub mod wrapper;
 
+mod lints;
 mod utils;
diff --git a/rust-cuda-derive/src/kernel/wrapper/generate/cpu_linker_macro/get_ptx_str.rs b/rust-cuda-derive/src/kernel/wrapper/generate/cpu_linker_macro/get_ptx_str.rs
@@ -5,6 +5,7 @@ use crate::kernel::utils::skip_kernel_compilation;
 
 use super::super::super::{DeclGenerics, FuncIdent, FunctionInputs, InputCudaType, KernelConfig};
 
+#[allow(clippy::too_many_arguments)]
 pub(super) fn quote_get_ptx_str(
     crate_path: &syn::Path,
     FuncIdent {
@@ -21,6 +22,7 @@ pub(super) fn quote_get_ptx_str(
     inputs: &FunctionInputs,
     func_params: &[syn::Ident],
     macro_type_ids: &[syn::Ident],
+    ptx_lint_levels: &TokenStream,
 ) -> TokenStream {
     let crate_name = match proc_macro::tracked_env::var("CARGO_CRATE_NAME") {
         Ok(crate_name) => crate_name.to_uppercase(),
@@ -80,7 +82,7 @@ pub(super) fn quote_get_ptx_str(
             #crate_path::host::link_kernel!{
                 #func_ident #func_ident_hash #args #crate_name #crate_manifest_dir #generic_start_token
                     #($#macro_type_ids),*
-                #generic_close_token
+                #generic_close_token #ptx_lint_levels
             }
 
             #matching_kernel_assert
diff --git a/rust-cuda-derive/src/kernel/wrapper/generate/cpu_linker_macro/mod.rs b/rust-cuda-derive/src/kernel/wrapper/generate/cpu_linker_macro/mod.rs
diff --git a/rust-cuda-derive/src/kernel/wrapper/mod.rs b/rust-cuda-derive/src/kernel/wrapper/mod.rs
