sm

hercules_cg/src/gpu.rs

@@ -155,7 +155,7 @@ pub fn gpu_codegen<W: Write>(
         (NodeID::new(pos), *data)
     };
 
-    let return_type_id = &typing[return_node_id.idx()];
+    let return_type_id = &typing[data_node_id.idx()];
     let return_type = &types[return_type_id.idx()];
     if return_type.is_array() || return_type.is_product() || return_type.is_summation() {
         let objects = &collection_objects.objects(data_node_id);
@@ -520,12 +520,13 @@ namespace cg = cooperative_groups;
 
     fn codegen_launch_code(&self, run_debug: bool, num_blocks: usize, num_threads: usize, dynamic_shared_offset: &str, w: &mut String) -> Result<(), Error> {
         write!(w, "
-int main() {{
-")?;
+int main(")?;
         // The following steps are for host-side C function arguments, but we also
         // need to pass arguments to kernel, so we keep track of the arguments here.
         let mut pass_args = String::new();
         if run_debug {
+            write!(w, ") {{
+")?;
             // The first set of parameters are dynamic constants.
             let mut first_param = true;
             for idx in 0..self.function.num_dynamic_constants {