3 files changed, 83 insertions, 1 deletions
diff --git a/ptx/src/pass/insert_explicit_load_store.rs b/ptx/src/pass/insert_explicit_load_store.rs
index 60c4a14..702f733 100644
--- a/ptx/src/pass/insert_explicit_load_store.rs
+++ b/ptx/src/pass/insert_explicit_load_store.rs
@@ -122,6 +122,13 @@ fn run_statement<'a, 'input>(
             result.push(Statement::Instruction(instruction));
             result.extend(visitor.post.drain(..).map(Statement::Instruction));
         }
+        Statement::PtrAccess(ptr_access) => {
+            let statement = Statement::PtrAccess(visitor.visit_ptr_access(ptr_access)?);
+            let statement = statement.visit_map(visitor)?;
+            result.extend(visitor.pre.drain(..).map(Statement::Instruction));
+            result.push(statement);
+            result.extend(visitor.post.drain(..).map(Statement::Instruction));
+        }
         s => {
             let new_statement = s.visit_map(visitor)?;
             result.extend(visitor.pre.drain(..).map(Statement::Instruction));
@@ -259,6 +266,41 @@ impl<'a, 'input> InsertMemSSAVisitor<'a, 'input> {
         Ok(ast::Instruction::Ld { data, arguments })
     }
 
+    fn visit_ptr_access(
+        &mut self,
+        ptr_access: PtrAccess<SpirvWord>,
+    ) -> Result<PtrAccess<SpirvWord>, TranslateError> {
+        let (old_space, new_space, name) = match self.variables.get(&ptr_access.ptr_src) {
+            Some(RemapAction::LDStSpaceChange {
+                old_space,
+                new_space,
+                name,
+            }) => (*old_space, *new_space, *name),
+            Some(RemapAction::PreLdPostSt { .. }) | None => return Ok(ptr_access),
+        };
+        if ptr_access.state_space != old_space {
+            return Err(error_mismatched_type());
+        }
+        // Propagate space changes in dst
+        let new_dst = self
+            .resolver
+            .register_unnamed(Some((ptr_access.underlying_type.clone(), new_space)));
+        self.variables.insert(
+            ptr_access.dst,
+            RemapAction::LDStSpaceChange {
+                old_space,
+                new_space,
+                name: new_dst,
+            },
+        );
+        Ok(PtrAccess {
+            ptr_src: name,
+            dst: new_dst,
+            state_space: new_space,
+            ..ptr_access
+        })
+    }
+
     fn visit_variable(&mut self, var: &mut ast::Variable<SpirvWord>) -> Result<(), TranslateError> {
         let old_space = match var.state_space {
             space @ (ptx_parser::StateSpace::Reg | ptx_parser::StateSpace::Param) => space,
diff --git a/ptx/src/pass/mod.rs b/ptx/src/pass/mod.rs
index ef131b4..c32cc39 100644
--- a/ptx/src/pass/mod.rs
+++ b/ptx/src/pass/mod.rs
@@ -22,6 +22,7 @@ mod normalize_identifiers2;
 mod normalize_predicates2;
 mod replace_instructions_with_function_calls;
 mod resolve_function_pointers;
+mod replace_known_functions;
 
 static ZLUDA_PTX_IMPL: &'static [u8] = include_bytes!("../../lib/zluda_ptx_impl.bc");
 const ZLUDA_PTX_PREFIX: &'static str = "__zluda_ptx_impl_";
@@ -42,9 +43,10 @@ pub fn to_llvm_module<'input>(ast: ast::Module<'input>) -> Result<Module, Transl
     let mut scoped_resolver = ScopedResolver::new(&mut flat_resolver);
     let sreg_map = SpecialRegistersMap2::new(&mut scoped_resolver)?;
     let directives = normalize_identifiers2::run(&mut scoped_resolver, ast.directives)?;
+    let directives = replace_known_functions::run(&flat_resolver, directives);
     let directives = normalize_predicates2::run(&mut flat_resolver, directives)?;
     let directives = resolve_function_pointers::run(directives)?;
-    let directives = fix_special_registers2::run(&mut flat_resolver, &sreg_map, directives)?;
+    let directives: Vec<Directive2<'_, ptx_parser::Instruction<ptx_parser::ParsedOperand<SpirvWord>>, ptx_parser::ParsedOperand<SpirvWord>>> = fix_special_registers2::run(&mut flat_resolver, &sreg_map, directives)?;
     let directives = expand_operands::run(&mut flat_resolver, directives)?;
     let directives = deparamize_functions::run(&mut flat_resolver, directives)?;
     let directives = insert_explicit_load_store::run(&mut flat_resolver, directives)?;
diff --git a/ptx/src/pass/replace_known_functions.rs b/ptx/src/pass/replace_known_functions.rs
new file mode 100644
index 0000000..56bb7e6
--- /dev/null
+++ b/ptx/src/pass/replace_known_functions.rs
@@ -0,0 +1,38 @@
+use super::{GlobalStringIdentResolver2, NormalizedDirective2, SpirvWord};
+
+pub(crate) fn run<'input>(
+    resolver: &GlobalStringIdentResolver2<'input>,
+    mut directives: Vec<NormalizedDirective2<'input>>,
+) -> Vec<NormalizedDirective2<'input>> {
+    for directive in directives.iter_mut() {
+        match directive {
+            NormalizedDirective2::Method(func) => {
+                func.import_as =
+                    replace_with_ptx_impl(resolver, &func.func_decl.name, func.import_as.take());
+            }
+            _ => {}
+        }
+    }
+    directives
+}
+
+fn replace_with_ptx_impl<'input>(
+    resolver: &GlobalStringIdentResolver2<'input>,
+    fn_name: &ptx_parser::MethodName<'input, SpirvWord>,
+    name: Option<String>,
+) -> Option<String> {
+    let known_names = ["__assertfail"];
+    match name {
+        Some(name) if known_names.contains(&&*name) => Some(format!("__zluda_ptx_impl_{}", name)),
+        Some(name) => Some(name),
+        None => match fn_name {
+            ptx_parser::MethodName::Func(name) => match resolver.ident_map.get(name) {
+                Some(super::IdentEntry {
+                    name: Some(name), ..
+                }) => Some(format!("__zluda_ptx_impl_{}", name)),
+                _ => None,
+            },
+            ptx_parser::MethodName::Kernel(..) => None,
+        },
+    }
+}