build2cmake: cxx-flags option for C++ compile flags for kernels (#171)

Author: danieldk

build2cmake/src/config/v2.rs

@@ -90,18 +90,21 @@ pub enum Kernel {
     Cuda {
         cuda_capabilities: Option<Vec<String>>,
         cuda_flags: Option<Vec<String>>,
+        cxx_flags: Option<Vec<String>>,
         depends: Vec<Dependencies>,
         include: Option<Vec<String>>,
         src: Vec<String>,
     },
     #[serde(rename_all = "kebab-case")]
     Metal {
+        cxx_flags: Option<Vec<String>>,
         depends: Vec<Dependencies>,
         include: Option<Vec<String>>,
         src: Vec<String>,
     },
     #[serde(rename_all = "kebab-case")]
     Rocm {
+        cxx_flags: Option<Vec<String>>,
         depends: Vec<Dependencies>,
         rocm_archs: Option<Vec<String>>,
         include: Option<Vec<String>>,
@@ -110,11 +113,19 @@ pub enum Kernel {
 }
 
 impl Kernel {
+    pub fn cxx_flags(&self) -> Option<&[String]> {
+        match self {
+            Kernel::Cuda { cxx_flags, .. }
+            | Kernel::Metal { cxx_flags, .. }
+            | Kernel::Rocm { cxx_flags, .. } => cxx_flags.as_deref(),
+        }
+    }
+
     pub fn include(&self) -> Option<&[String]> {
         match self {
-            Kernel::Cuda { include, .. } => include.as_deref(),
-            Kernel::Metal { include, .. } => include.as_deref(),
-            Kernel::Rocm { include, .. } => include.as_deref(),
+            Kernel::Cuda { include, .. }
+            | Kernel::Metal { include, .. }
+            | Kernel::Rocm { include, .. } => include.as_deref(),
         }
     }
 
@@ -128,17 +139,15 @@ impl Kernel {
 
     pub fn depends(&self) -> &[Dependencies] {
         match self {
-            Kernel::Cuda { depends, .. } => depends,
-            Kernel::Metal { depends, .. } => depends,
-            Kernel::Rocm { depends, .. } => depends,
+            Kernel::Cuda { depends, .. }
+            | Kernel::Metal { depends, .. }
+            | Kernel::Rocm { depends, .. } => depends,
         }
     }
 
     pub fn src(&self) -> &[String] {
         match self {
-            Kernel::Cuda { src, .. } => src,
-            Kernel::Metal { src, .. } => src,
-            Kernel::Rocm { src, .. } => src,
+            Kernel::Cuda { src, .. } | Kernel::Metal { src, .. } | Kernel::Rocm { src, .. } => src,
         }
     }
 }
@@ -233,6 +242,7 @@ fn convert_kernels(v1_kernels: HashMap<String, v1::Kernel>) -> Result<HashMap<St
             kernels.insert(
                 format!("{name}_rocm"),
                 Kernel::Rocm {
+                    cxx_flags: None,
                     rocm_archs: kernel.rocm_archs,
                     depends: kernel.depends.clone(),
                     include: kernel.include.clone(),
@@ -246,6 +256,7 @@ fn convert_kernels(v1_kernels: HashMap<String, v1::Kernel>) -> Result<HashMap<St
             Kernel::Cuda {
                 cuda_capabilities: kernel.cuda_capabilities,
                 cuda_flags: None,
+                cxx_flags: None,
                 depends: kernel.depends,
                 include: kernel.include,
                 src: kernel.src,

build2cmake/src/templates/cuda/kernel.cmake

@@ -32,6 +32,16 @@ if(GPU_LANG STREQUAL "CUDA")
   endforeach()
   {% endif %}
 
+  {% if cxx_flags %}
+  foreach(_KERNEL_SRC {{'${' + kernel_name + '_SRC}'}})
+    set_property(
+      SOURCE ${_KERNEL_SRC}
+      APPEND PROPERTY
+      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CXX>:{{ cxx_flags }}>"
+    )
+  endforeach()
+  {% endif %}
+
   list(APPEND SRC {{'"${' + kernel_name + '_SRC}"'}})
 {% if supports_hipify %}
 elseif(GPU_LANG STREQUAL "HIP")

build2cmake/src/templates/metal/kernel.cmake

@@ -22,10 +22,20 @@ set_source_files_properties(
   PROPERTIES INCLUDE_DIRECTORIES "{{ includes }}")
 {% endif %}
 
+{% if cxx_flags %}
+foreach(_KERNEL_SRC {{'${' + kernel_name + '_CPP_SRC}'}})
+  set_property(
+    SOURCE ${_KERNEL_SRC}
+    APPEND PROPERTY
+    COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CXX>:{{ cxx_flags }}>"
+  )
+endforeach()
+{% endif %}
+
 # Add C++ sources to main source list
 list(APPEND SRC {{'"${' + kernel_name + '_CPP_SRC}"'}})
 
 # Keep track of Metal sources for later compilation
 if({{kernel_name}}_METAL_SRC)
   list(APPEND ALL_METAL_SOURCES {{'"${' + kernel_name + '_METAL_SRC}"'}})
-endif()
+endif()

build2cmake/src/torch/cuda.rs

@@ -314,6 +314,7 @@ pub fn render_kernel(
             context! {
                 cuda_capabilities => cuda_capabilities,
                 cuda_flags => cuda_flags.map(|flags| flags.join(";")),
+                cxx_flags => kernel.cxx_flags().map(|flags| flags.join(";")),
                 rocm_archs => rocm_archs,
                 includes => kernel.include().map(prefix_and_join_includes),
                 kernel_name => kernel_name,

build2cmake/src/torch/metal.rs

@@ -158,6 +158,7 @@ pub fn render_kernel(
         .wrap_err("Cannot get kernel template")?
         .render_to_write(
             context! {
+                cxx_flags => kernel.cxx_flags().map(|flags| flags.join(";")),
                 includes => kernel.include().map(prefix_and_join_includes),
                 kernel_name => kernel_name,
                 sources => sources,