Descend Arena Integration

Cargo.toml

@@ -6,15 +6,46 @@ edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[profile.samply]
+inherits = "release"
+debug = true
+
+[[bin]]
+name = "descendc"
+path = "src/main.rs"
+
 [dependencies.peg]
 version = "0.8.0"
 
+[dependencies.bumpalo]
+version="3.17.0"
+features=["collections", "boxed"]
+
 [dependencies.annotate-snippets]
 version = "0.9.0"
 features = ["color"]
 
 [dependencies.descend_derive]
 path = "./descend_derive"
 
+[dependencies.clap]
+version = "4.3"
+features = ["derive"]
+
+[dependencies.which]
+version = "7.0.2"
+
+[dependencies.log]
+version = "0.4.27"
+
+[dependencies.env_logger]
+version = "0.11.7"
+
+[dependencies.predicates]
+version= "3.1.3"
+
+[dependencies.assert_cmd]
+version = "2.0.16"
+
 [workspace]
-members = ["descend_derive"]
+members = ["descend_derive"]

README.md

@@ -110,7 +110,7 @@ codegen
 
 cuda-examples/
 ---------------------
-* Contains handwritte or generated CUDA programs
+* Contains handwritten or generated CUDA programs
 * Contains `descend.cuh`; the header file which is required in order to compile Descend programs,
   that were translated to CUDA, with `nvcc` (contains for example the implementation of `exec`)
 

descend_derive/src/lib.rs

@@ -1,14 +1,161 @@
-use std::collections::HashSet;
-
 use proc_macro::TokenStream;
 use quote::quote;
+use std::collections::HashSet;
 use syn::{
     parse::{Parse, ParseStream},
     parse_macro_input,
     punctuated::Punctuated,
     Fields, FieldsNamed, Ident, ItemStruct, Token, Type, TypeReference,
 };
 
+struct Args {
+    vars: HashSet<Ident>,
+}
+
+impl Parse for Args {
+    fn parse(input: ParseStream) -> syn::parse::Result<Self> {
+        let vars = Punctuated::<Ident, Token![,]>::parse_terminated(input)?;
+        Ok(Args {
+            vars: vars.into_iter().collect(),
+        })
+    }
+}
+
+const IGNORE_ATTR_NAME: &str = "span_derive_ignore";
+
+#[proc_macro_attribute]
+pub fn span_derive(attr: TokenStream, input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as ItemStruct);
+    let args = parse_macro_input!(attr as Args);
+    if args.vars.is_empty() {
+        panic!("Need at least one attribute in span_derive!");
+    }
+
+    let original_name = &input.ident;
+    let original_generics = input.generics.clone();
+    let (impl_generics, ty_generics, where_clause) = original_generics.split_for_impl();
+
+    let original_fields = match &input.fields {
+        syn::Fields::Named(fields) => fields,
+        _ => unreachable!(),
+    };
+
+    let mut new_fields = original_fields.clone();
+    for field in new_fields.named.iter_mut() {
+        field
+            .attrs
+            .retain(|attr| !attr.path.is_ident(IGNORE_ATTR_NAME));
+    }
+
+    let mut original = input.clone();
+    original.fields = Fields::Named(new_fields);
+
+    let mut helper = input.clone();
+
+    let mut helper_generics = original_generics.clone();
+    let mut new_params = Punctuated::new();
+    new_params.push(syn::parse_quote!('helper));
+    for param in original_generics.params.iter() {
+        new_params.push(param.clone());
+    }
+    helper_generics.params = new_params;
+    helper.generics = helper_generics.clone();
+
+    let mut helper_fields = FieldsNamed {
+        brace_token: original_fields.brace_token.clone(),
+        named: Punctuated::new(),
+    };
+    for mut field in original_fields.named.clone().into_iter() {
+        if field
+            .attrs
+            .iter()
+            .any(|attr| attr.path.is_ident(IGNORE_ATTR_NAME))
+        {
+            continue;
+        }
+        let mut typ: TypeReference = syn::parse_str("&'helper i64").unwrap();
+        typ.elem = Box::new(field.ty);
+        field.ty = Type::Reference(typ);
+        helper_fields.named.push(field);
+    }
+    helper.fields = Fields::Named(helper_fields);
+
+    let helper_name = Ident::new(
+        &format!("__{}SpanHelper", original_name.to_string()),
+        input.ident.span(),
+    );
+    helper.ident = helper_name.clone();
+    helper.attrs.clear();
+
+    let into_fields = helper
+        .fields
+        .iter()
+        .map(|field| field.ident.as_ref().unwrap().clone())
+        .collect::<Vec<_>>();
+
+    let mut derive_args: Vec<_> = args.vars.iter().collect();
+    derive_args.sort_by(|a, b| a.to_string().cmp(&b.to_string()));
+
+    let (helper_impl_generics, helper_ty_generics, helper_where_clause) =
+        helper_generics.split_for_impl();
+
+    let mut output = quote! {
+        #original
+
+        #[derive(#(#derive_args),*)]
+        #helper
+
+        impl #helper_impl_generics From<&'helper #original_name #ty_generics>
+            for #helper_name #helper_ty_generics #helper_where_clause {
+            fn from(orig: &'helper #original_name #ty_generics) -> Self {
+                #helper_name {
+                    #(#into_fields: &orig.#into_fields),*
+                }
+            }
+        }
+    };
+
+    for derive_arg in args.vars.iter() {
+        match &derive_arg.to_string() as &str {
+            "PartialEq" => {
+                output.extend(quote! {
+                    impl #impl_generics ::core::cmp::PartialEq for #original_name #ty_generics #where_clause {
+                        fn eq(&self, other: &Self) -> bool {
+                            let helper = #helper_name::from(self);
+                            let helper_other = #helper_name::from(other);
+                            helper == helper_other
+                        }
+                    }
+                });
+            }
+            "Eq" => {
+                output.extend(quote! {
+                    impl #impl_generics ::core::cmp::Eq for #original_name #ty_generics #where_clause {}
+                });
+            }
+            "Hash" => {
+                output.extend(quote! {
+                    impl #impl_generics ::core::hash::Hash for #original_name #ty_generics #where_clause {
+                        fn hash<H: ::core::hash::Hasher>(&self, state: &mut H) {
+                            let helper = #helper_name::from(self);
+                            helper.hash(state);
+                        }
+                    }
+                });
+            }
+            other => panic!("span_derive not implemented for {}", other),
+        }
+    }
+
+    eprintln!(
+        "--- Generated code ---\n{}\n------------------------",
+        output
+    );
+
+    TokenStream::from(output)
+}
+
+/*** OLD WAY
 // Copy paste from syn example
 struct Args {
     vars: HashSet<Ident>,
@@ -148,3 +295,4 @@ pub fn span_derive(attr: TokenStream, input: TokenStream) -> TokenStream {
 
     TokenStream::from(output)
 }
+*/

profile_memory.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+INPUT="examples/infer/matmul.desc"
+
+cargo build --profile samply
+
+samply record -- ./target/samply/descendc emit "$INPUT"