Start refactoring SPIRV module generation in preparation for support of functions

notcuda/src/impl/module.rs

@@ -14,10 +14,7 @@ pub enum ModuleCompileError<'a> {
 }
 
 impl<'a> ModuleCompileError<'a> {
-    pub fn get_build_log(&self) {
-
-    }
-
+    pub fn get_build_log(&self) {}
 }
 
 impl<'a> From<ptx::SpirvError> for ModuleCompileError<'a> {

ptx/src/ast.rs

@@ -48,24 +48,25 @@ impl<
 
 pub struct Module<'a> {
     pub version: (u8, u8),
-    pub functions: Vec<Function<'a>>,
+    pub functions: Vec<ParsedFunction<'a>>,
 }
 
-pub enum FunctionReturn<'a> {
-    Func(Vec<Argument<'a>>),
-    Kernel,
+pub enum FunctionHeader<'a, P: ArgParams> {
+    Func(Vec<Argument<P>>, P::ID),
+    Kernel(&'a str),
 }
 
-pub struct Function<'a> {
-    pub func_directive: FunctionReturn<'a>,
-    pub name: &'a str,
-    pub args: Vec<Argument<'a>>,
-    pub body: Option<Vec<Statement<ParsedArgParams<'a>>>>,
+pub struct Function<'a, P: ArgParams, S> {
+    pub func_directive: FunctionHeader<'a, P>,
+    pub args: Vec<Argument<P>>,
+    pub body: Option<Vec<S>>,
 }
 
+pub type ParsedFunction<'a> = Function<'a, ParsedArgParams<'a>, Statement<ParsedArgParams<'a>>>;
+
 #[derive(Default)]
-pub struct Argument<'a> {
-    pub name: &'a str,
+pub struct Argument<P: ArgParams> {
+    pub name: P::ID,
     pub a_type: ScalarType,
     pub length: u32,
 }
@@ -231,7 +232,7 @@ pub struct CallData {
 
 pub struct AbsDetails {
     pub flush_to_zero: bool,
-    pub typ: ScalarType
+    pub typ: ScalarType,
 }
 
 pub struct ArgCall<P: ArgParams> {

ptx/src/ptx.lalrpop

@@ -189,7 +189,7 @@ TargetSpecifier = {
     "map_f64_to_f32"
 };
 
-Directive: Option<ast::Function<'input>> = {
+Directive: Option<ast::Function<'input, ast::ParsedArgParams<'input>, ast::Statement<ast::ParsedArgParams<'input>>>> = {
     AddressSize => None,
     <f:Function> => Some(f),
     File => None,
@@ -200,10 +200,9 @@ AddressSize = {
     ".address_size" Num
 };
 
-Function: ast::Function<'input>  = {
+Function: ast::Function<'input, ast::ParsedArgParams<'input>, ast::Statement<ast::ParsedArgParams<'input>>>  = {
     LinkingDirective*
-    <func_directive:FunctionReturn>
-    <name:ExtendedID>
+    <func_directive:FunctionHeader>
     <args:Arguments>
     <body:FunctionBody> =>  ast::Function{<>}
 };
@@ -214,17 +213,17 @@ LinkingDirective = {
     ".weak"
 };
 
-FunctionReturn: ast::FunctionReturn<'input> = {
-    ".entry" => ast::FunctionReturn::Kernel,
-    ".func" <args:Arguments?> => ast::FunctionReturn::Func(args.unwrap_or_else(|| Vec::new()))
+FunctionHeader: ast::FunctionHeader<'input, ast::ParsedArgParams<'input>> = {
+    ".entry" <name:ExtendedID> => ast::FunctionHeader::Kernel(name),
+    ".func" <args:Arguments?> <name:ExtendedID> => ast::FunctionHeader::Func(args.unwrap_or_else(|| Vec::new()), name)
 };
 
-Arguments: Vec<ast::Argument<'input>> = {
+Arguments: Vec<ast::Argument<ast::ParsedArgParams<'input>>> = {
     "(" <args:Comma<FunctionInput>> ")" => args
 }
 
 // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parameter-state-space
-FunctionInput: ast::Argument<'input> = {
+FunctionInput: ast::Argument<ast::ParsedArgParams<'input>> = {
     ".param" <_type:ScalarType> <name:ExtendedID> => { 
         ast::Argument {a_type: _type, name: name, length: 1 }
     },

ptx/src/test/spirv_run/mod.rs

@@ -52,6 +52,7 @@ test_ptx!(cvt_sat_s_u, [-1i32], [0i32]);
 test_ptx!(cvta, [3.0f32], [3.0f32]);
 test_ptx!(block, [1u64], [2u64]);
 test_ptx!(local_align, [1u64], [1u64]);
+test_ptx!(call, [1u64], [2u64]);
 
 struct DisplayError<T: Debug> {
     err: T,

ptx/src/translate.rs

@@ -154,7 +154,7 @@ impl TypeWordMap {
     }
 }
 
-pub fn to_spirv_module(ast: ast::Module) -> Result<dr::Module, dr::Error> {
+pub fn to_spirv_module<'a>(ast: ast::Module<'a>) -> Result<dr::Module, dr::Error> {
     let mut builder = dr::Builder::new();
     // https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html#_a_id_logicallayout_a_logical_layout_of_a_module
     builder.set_version(1, 3);
@@ -163,13 +163,21 @@ pub fn to_spirv_module(ast: ast::Module) -> Result<dr::Module, dr::Error> {
     let opencl_id = emit_opencl_import(&mut builder);
     emit_memory_model(&mut builder);
     let mut map = TypeWordMap::new(&mut builder);
-    for f in ast.functions {
-        emit_function(&mut builder, &mut map, opencl_id, f)?;
+    let mut id_defs = GlobalStringIdResolver::new(builder.id());
+    let ssa_functions = ast
+        .functions
+        .into_iter()
+        .map(|f| to_ssa_function(&mut id_defs, opencl_id, f))
+        .collect::<Vec<_>>();
+    for f in ssa_functions {
+        emit_function_args(&mut builder, &mut map, &*f.args);
+        emit_function_body_ops(&mut builder, &mut map, opencl_id, &f.body)?;
+        builder.end_function()?;
     }
     Ok(builder.module())
 }
 
-pub fn to_spirv(ast: ast::Module) -> Result<Vec<u32>, dr::Error> {
+pub fn to_spirv<'a>(ast: ast::Module<'a>) -> Result<Vec<u32>, dr::Error> {
     let module = to_spirv_module(ast)?;
     Ok(module.assemble())
 }
@@ -196,28 +204,28 @@ fn emit_memory_model(builder: &mut dr::Builder) {
     );
 }
 
-fn emit_function<'a>(
-    builder: &mut dr::Builder,
-    map: &mut TypeWordMap,
+fn to_ssa_function<'a>(
+    id_defs: &mut GlobalStringIdResolver<'a>,
     opencl_id: spirv::Word,
-    f: ast::Function<'a>,
-) -> Result<spirv::Word, rspirv::dr::Error> {
-    let func_type = get_function_type(builder, map, &f.args);
-    let func_id =
-        builder.begin_function(map.void(), None, spirv::FunctionControl::NONE, func_type)?;
-    match f.func_directive {
-        ast::FunctionReturn::Kernel => {
-            builder.entry_point(spirv::ExecutionModel::Kernel, func_id, f.name, &[])
+    f: ast::ParsedFunction<'a>,
+) -> ExpandedFunction<'a> {
+    let ids_start = id_defs.current_id();
+    let fn_resolver = FnStringIdResolver::new(id_defs);
+    let f_header = match f.func_directive {
+        ast::FunctionHeader::Kernel(name) => todo!(),
+        ast::FunctionHeader::Func(ret_params, name) => todo!(),
+    };
+    let f_args = todo!();
+    let f_body = Some(to_ssa(
+        fn_resolver,
+        &f.args,
+        f.body.unwrap_or_else(|| todo!()),
+    ));
+    ExpandedFunction {
+        func_directive: f_header,
+        args: f_args,
+        body: f_body,
     }
-        _ => todo!(),
-    }
-    let (mut func_body, unique_ids) = to_ssa(&f.args, f.body.unwrap_or_else(|| todo!()));
-    let id_offset = builder.reserve_ids(unique_ids);
-    emit_function_args(builder, id_offset, map, &f.args);
-    func_body = apply_id_offset(func_body, id_offset);
-    emit_function_body_ops(builder, map, opencl_id, &func_body)?;
-    builder.end_function()?;
-    Ok(func_id)
 }
 
 fn apply_id_offset(func_body: Vec<ExpandedStatement>, id_offset: u32) -> Vec<ExpandedStatement> {
@@ -228,16 +236,19 @@ fn apply_id_offset(func_body: Vec<ExpandedStatement>, id_offset: u32) -> Vec<Exp
 }
 
 fn to_ssa<'a, 'b>(
-    f_args: &'b [ast::Argument<'a>],
+    mut id_defs: FnStringIdResolver<'a, 'b>,
+    f_args: &'b [ast::Argument<ast::ParsedArgParams<'a>>],
     f_body: Vec<ast::Statement<ast::ParsedArgParams<'a>>>,
-) -> (Vec<ExpandedStatement>, spirv::Word) {
-    let (normalized_ids, mut id_def) = normalize_identifiers(&f_args, f_body);
-    let normalized_statements = normalize_predicates(normalized_ids, &mut id_def);
-    let ssa_statements = insert_mem_ssa_statements(normalized_statements, &mut id_def);
-    let expanded_statements = expand_arguments(ssa_statements, &mut id_def);
-    let expanded_statements = insert_implicit_conversions(expanded_statements, &mut id_def);
-    let labeled_statements = normalize_labels(expanded_statements, &mut id_def);
-    (labeled_statements, id_def.ids_count())
+) -> Vec<ExpandedStatement> {
+    let normalized_ids = normalize_identifiers(&mut id_defs, &f_args, f_body);
+    let mut numeric_id_defs = id_defs.finish();
+    let normalized_statements = normalize_predicates(normalized_ids, &mut numeric_id_defs);
+    let ssa_statements = insert_mem_ssa_statements(normalized_statements, &mut numeric_id_defs);
+    let expanded_statements = expand_arguments(ssa_statements, &mut numeric_id_defs);
+    let expanded_statements =
+        insert_implicit_conversions(expanded_statements, &mut numeric_id_defs);
+    let labeled_statements = normalize_labels(expanded_statements, &mut numeric_id_defs);
+    labeled_statements
 }
 
 fn normalize_labels(
@@ -391,9 +402,9 @@ fn insert_mem_ssa_statements(
     result
 }
 
-fn expand_arguments(
+fn expand_arguments<'a, 'b, 'c>(
     func: Vec<NormalizedStatement>,
-    id_def: &mut NumericIdResolver,
+    id_def: &'c mut NumericIdResolver<'a, 'b>,
 ) -> Vec<ExpandedStatement> {
     let mut result = Vec::with_capacity(func.len());
     for s in func {
@@ -416,18 +427,23 @@ fn expand_arguments(
     result
 }
 
-struct FlattenArguments<'a> {
-    func: &'a mut Vec<ExpandedStatement>,
-    id_def: &'a mut NumericIdResolver,
+struct FlattenArguments<'a, 'b, 'c> {
+    func: &'c mut Vec<ExpandedStatement>,
+    id_def: &'c mut NumericIdResolver<'a, 'b>,
 }
 
-impl<'a> FlattenArguments<'a> {
-    fn new(func: &'a mut Vec<ExpandedStatement>, id_def: &'a mut NumericIdResolver) -> Self {
+impl<'a, 'b, 'c> FlattenArguments<'a, 'b, 'c> {
+    fn new(
+        func: &'c mut Vec<ExpandedStatement>,
+        id_def: &'c mut NumericIdResolver<'a, 'b>,
+    ) -> Self {
         FlattenArguments { func, id_def }
     }
 }
 
-impl<'a> ArgumentMapVisitor<NormalizedArgParams, ExpandedArgParams> for FlattenArguments<'a> {
+impl<'a, 'b, 'c> ArgumentMapVisitor<NormalizedArgParams, ExpandedArgParams>
+    for FlattenArguments<'a, 'b, 'c>
+{
     fn dst_variable(&mut self, desc: ArgumentDescriptor<spirv::Word>) -> spirv::Word {
         desc.op
     }
@@ -577,18 +593,17 @@ fn insert_implicit_conversions(
 fn get_function_type(
     builder: &mut dr::Builder,
     map: &mut TypeWordMap,
-    args: &[ast::Argument],
+    args: &[ast::Argument<ast::ParsedArgParams>],
 ) -> spirv::Word {
     map.get_or_add_fn(builder, args.iter().map(|arg| SpirvType::from(arg.a_type)))
 }
 
 fn emit_function_args(
     builder: &mut dr::Builder,
-    id_offset: spirv::Word,
     map: &mut TypeWordMap,
-    args: &[ast::Argument],
+    args: &[ast::Argument<ExpandedArgParams>],
 ) {
-    let mut id = id_offset;
+    let mut id = todo!();
     for arg in args {
         let result_type = map.get_or_add_scalar(builder, arg.a_type);
         let inst = dr::Instruction::new(
@@ -606,9 +621,9 @@ fn emit_function_body_ops(
     builder: &mut dr::Builder,
     map: &mut TypeWordMap,
     opencl: spirv::Word,
-    func: &[ExpandedStatement],
+    func: &Option<Vec<ExpandedStatement>>,
 ) -> Result<(), dr::Error> {
-    for s in func {
+    for s in func.as_ref().unwrap() {
         match s {
             Statement::Label(id) => {
                 if builder.block.is_some() {
@@ -1079,10 +1094,10 @@ fn emit_implicit_conversion(
 
 // TODO: support scopes
 fn normalize_identifiers<'a, 'b>(
-    args: &'b [ast::Argument<'a>],
+    id_defs: &mut FnStringIdResolver<'a, 'b>,
+    args: &[ast::Argument<ast::ParsedArgParams<'a>>],
     func: Vec<ast::Statement<ast::ParsedArgParams<'a>>>,
-) -> (Vec<ast::Statement<NormalizedArgParams>>, NumericIdResolver) {
-    let mut id_defs = StringIdResolver::new();
+) -> Vec<ast::Statement<NormalizedArgParams>> {
     for arg in args {
         id_defs.add_def(arg.name, Some(ast::Type::Scalar(arg.a_type)));
     }
@@ -1096,13 +1111,13 @@ fn normalize_identifiers<'a, 'b>(
     }
     let mut result = Vec::new();
     for s in func {
-        expand_map_variables(&mut id_defs, &mut result, s);
+        expand_map_variables(id_defs, &mut result, s);
     }
-    (result, id_defs.finish())
+    result
 }
 
-fn expand_map_variables<'a>(
-    id_defs: &mut StringIdResolver<'a>,
+fn expand_map_variables<'a, 'b>(
+    id_defs: &mut FnStringIdResolver<'a, 'b>,
     result: &mut Vec<ast::Statement<NormalizedArgParams>>,
     s: ast::Statement<ast::ParsedArgParams<'a>>,
 ) {
@@ -1145,24 +1160,53 @@ fn expand_map_variables<'a>(
     }
 }
 
-struct StringIdResolver<'a> {
+struct GlobalStringIdResolver<'a> {
     current_id: spirv::Word,
+    variables: HashMap<Cow<'a, str>, spirv::Word>,
+}
+
+impl<'a> GlobalStringIdResolver<'a> {
+    fn new(start_id: spirv::Word) -> Self {
+        Self {
+            current_id: start_id,
+            variables: HashMap::new(),
+        }
+    }
+
+    fn add_def(&mut self, id: &'a str) -> spirv::Word {
+        let numeric_id = self.current_id;
+        self.variables.insert(Cow::Borrowed(id), numeric_id);
+        self.current_id += 1;
+        numeric_id
+    }
+
+    fn reserve_id(&mut self) {
+        self.current_id += 1;
+    }
+
+    fn current_id(&self) -> spirv::Word {
+        self.current_id
+    }
+}
+
+struct FnStringIdResolver<'a, 'b> {
+    global: &'b mut GlobalStringIdResolver<'a>,
     variables: Vec<HashMap<Cow<'a, str>, spirv::Word>>,
     type_check: HashMap<u32, ast::Type>,
 }
 
-impl<'a> StringIdResolver<'a> {
-    fn new() -> Self {
-        StringIdResolver {
-            current_id: 0u32,
+impl<'a, 'b> FnStringIdResolver<'a, 'b> {
+    fn new(global: &'b mut GlobalStringIdResolver<'a>) -> Self {
+        Self {
+            global: global,
             variables: vec![HashMap::new(); 1],
             type_check: HashMap::new(),
         }
     }
 
-    fn finish(self) -> NumericIdResolver {
+    fn finish(self) -> NumericIdResolver<'a, 'b> {
         NumericIdResolver {
-            current_id: self.current_id,
+            global: self.global,
             type_check: self.type_check,
         }
     }
@@ -1175,18 +1219,18 @@ impl<'a> StringIdResolver<'a> {
         self.variables.pop();
     }
 
-    fn get_id(&self, id: &'a str) -> spirv::Word {
+    fn get_id(&self, id: &str) -> spirv::Word {
         for scope in self.variables.iter().rev() {
             match scope.get(id) {
                 Some(id) => return *id,
                 None => continue,
             }
         }
-        panic!()
+        self.global.variables[id]
     }
 
     fn add_def(&mut self, id: &'a str, typ: Option<ast::Type>) -> spirv::Word {
-        let numeric_id = self.current_id;
+        let numeric_id = self.global.current_id;
         self.variables
             .last_mut()
             .unwrap()
@@ -1194,7 +1238,7 @@ impl<'a> StringIdResolver<'a> {
         if let Some(typ) = typ {
             self.type_check.insert(numeric_id, typ);
         }
-        self.current_id += 1;
+        self.global.current_id += 1;
         numeric_id
     }
 
@@ -1205,7 +1249,7 @@ impl<'a> StringIdResolver<'a> {
         count: u32,
         typ: ast::Type,
     ) -> impl Iterator<Item = spirv::Word> {
-        let numeric_id = self.current_id;
+        let numeric_id = self.global.current_id;
         for i in 0..count {
             self.variables
                 .last_mut()
@@ -1213,33 +1257,29 @@ impl<'a> StringIdResolver<'a> {
                 .insert(Cow::Owned(format!("{}{}", base_id, i)), numeric_id + i);
             self.type_check.insert(numeric_id + i, typ);
         }
-        self.current_id += count;
+        self.global.current_id += count;
         (0..count).into_iter().map(move |i| i + numeric_id)
     }
 }
 
-struct NumericIdResolver {
-    current_id: spirv::Word,
+struct NumericIdResolver<'a, 'b> {
+    global: &'b mut GlobalStringIdResolver<'a>,
     type_check: HashMap<u32, ast::Type>,
 }
 
-impl NumericIdResolver {
+impl<'a, 'b> NumericIdResolver<'a, 'b> {
     fn get_type(&self, id: spirv::Word) -> ast::Type {
         self.type_check[&id]
     }
 
     fn new_id(&mut self, typ: Option<ast::Type>) -> spirv::Word {
-        let new_id = self.current_id;
+        let new_id = self.global.current_id;
         if let Some(typ) = typ {
             self.type_check.insert(new_id, typ);
         }
-        self.current_id += 1;
+        self.global.current_id += 1;
         new_id
     }
-
-    fn ids_count(&self) -> spirv::Word {
-        self.current_id
-    }
 }
 
 enum Statement<I> {
@@ -1284,6 +1324,7 @@ impl ast::ArgParams for NormalizedArgParams {
 
 enum ExpandedArgParams {}
 type ExpandedStatement = Statement<ast::Instruction<ExpandedArgParams>>;
+type ExpandedFunction<'a> = ast::Function<'a, ExpandedArgParams, ExpandedStatement>;
 
 impl ast::ArgParams for ExpandedArgParams {
     type ID = spirv::Word;