Implement vmprof support for aarch64

7 files changed, 76 insertions, 11 deletions

diff --git a/rpython/jit/backend/aarch64/assembler.py b/rpython/jit/backend/aarch64/assembler.py
index 391716a75c..050710f1fa 100644
--- a/rpython/jit/backend/aarch64/assembler.py
+++ b/rpython/jit/backend/aarch64/assembler.py
@@ -44,6 +44,9 @@ class AssemblerARM64(ResOpAssembler):
             assert len(set(inputargs)) == len(inputargs)
 
         self.setup(looptoken)
+        if self.cpu.HAS_CODEMAP:
+            self.codemap_builder.enter_portal_frame(jd_id, unique_id,
+                                                    self.mc.get_relative_pos())
 
         frame_info = self.datablockwrapper.malloc_aligned(
             jitframe.JITFRAMEINFO_SIZE, alignment=WORD)
@@ -124,7 +127,10 @@ class AssemblerARM64(ResOpAssembler):
             assert len(set(inputargs)) == len(inputargs)
 
         self.setup(original_loop_token)
-        #self.codemap.inherit_code_from_position(faildescr.adr_jump_offset)
+        if self.cpu.HAS_CODEMAP:
+            self.codemap_builder.inherit_code_from_position(
+                faildescr.adr_jump_offset)
+
         descr_number = compute_unique_id(faildescr)
         if log:
             operations = self._inject_debugging_code(faildescr, operations,
@@ -1019,8 +1025,8 @@ class AssemblerARM64(ResOpAssembler):
         size = self.mc.get_relative_pos() 
         res = self.mc.materialize(self.cpu, allblocks,
                                    self.cpu.gc_ll_descr.gcrootmap)
-        #self.cpu.codemap.register_codemap(
-        #    self.codemap.get_final_bytecode(res, size))
+        self.cpu.codemap.register_codemap(
+            self.codemap_builder.get_final_bytecode(res, size))
         return res
 
     def patch_trace(self, faildescr, looptoken, bridge_addr, regalloc):
@@ -1086,13 +1092,16 @@ class AssemblerARM64(ResOpAssembler):
             pmc.B_ofs_cond(self.mc.currpos() - pos, c.LS)
 
     def _call_header(self):
-        stack_size = (len(r.callee_saved_registers) + 4) * WORD
+        stack_size = (len(r.callee_saved_registers) + 8) * WORD
         self.mc.STP_rr_preindex(r.lr.value, r.fp.value, r.sp.value, -stack_size)
         for i in range(0, len(r.callee_saved_registers), 2):
             self.mc.STP_rri(r.callee_saved_registers[i].value,
                             r.callee_saved_registers[i + 1].value,
                             r.sp.value,
-                            (i + 4) * WORD)
+                            (i + 8) * WORD)
+
+        if self.cpu.translate_support_code:
+            self._call_header_vmprof()
         
         self.saved_threadlocal_addr = 3 * WORD   # at offset 3 from location 'sp'
         self.mc.STR_ri(r.x1.value, r.sp.value, 3 * WORD)
@@ -1100,10 +1109,33 @@ class AssemblerARM64(ResOpAssembler):
         # set fp to point to the JITFRAME, passed in argument 'x0'
         self.mc.MOV_rr(r.fp.value, r.x0.value)
         #
+
         gcrootmap = self.cpu.gc_ll_descr.gcrootmap
         if gcrootmap and gcrootmap.is_shadow_stack:
             self.gen_shadowstack_header(gcrootmap)
 
+    def _call_header_vmprof(self):
+        # this uses values 0, 1 and 2 on stack as vmprof next
+        from rpython.rlib.rvmprof.rvmprof import cintf, VMPROF_JITTED_TAG
+
+        # tloc = address of pypy_threadlocal_s
+        tloc = r.x1
+        # ip0 = current value of vmprof_tl_stack
+        offset = cintf.vmprof_tl_stack.getoffset()
+        self.mc.LDR_ri(r.ip0.value, tloc.value, offset)
+        # stack->next = old
+        self.mc.STR_ri(r.ip0.value, r.sp.value, 4 * WORD)
+        # stack->value = my sp
+        self.mc.ADD_ri(r.ip1.value, r.sp.value, 0)
+        self.mc.STR_ri(r.ip1.value, r.sp.value, (4 + 1) * WORD)
+        # stack->kind = VMPROF_JITTED_TAG
+        self.mc.gen_load_int(r.ip0.value, VMPROF_JITTED_TAG)
+        self.mc.STR_ri(r.ip0.value, r.sp.value, (4 + 2) * WORD)
+        # save in vmprof_tl_stack the new eax
+        self.mc.ADD_ri(r.ip0.value, r.sp.value, 4 * WORD)
+        self.mc.STR_ri(r.ip0.value, tloc.value, offset)
+
+
     def _assemble(self, regalloc, inputargs, operations):
         #self.guard_success_cc = c.cond_none
         regalloc.compute_hint_frame_locations(operations)
@@ -1339,20 +1371,32 @@ class AssemblerARM64(ResOpAssembler):
         if gcrootmap and gcrootmap.is_shadow_stack:
             self.gen_footer_shadowstack(gcrootmap, mc)
 
+        if self.cpu.translate_support_code:
+            self._call_footer_vmprof(mc)
         # pop all callee saved registers
 
-        stack_size = (len(r.callee_saved_registers) + 4) * WORD
+        stack_size = (len(r.callee_saved_registers) + 8) * WORD
 
         for i in range(0, len(r.callee_saved_registers), 2):
             mc.LDP_rri(r.callee_saved_registers[i].value,
                             r.callee_saved_registers[i + 1].value,
                             r.sp.value,
-                            (i + 4) * WORD)
+                            (i + 8) * WORD)
         mc.LDP_rr_postindex(r.lr.value, r.fp.value, r.sp.value, stack_size)
 
 
         mc.RET_r(r.lr.value)
 
+    def _call_footer_vmprof(self, mc):
+        from rpython.rlib.rvmprof.rvmprof import cintf
+        # ip0 = address of pypy_threadlocal_s
+        mc.LDR_ri(r.ip0.value, r.sp.value, 3 * WORD)
+        # ip1 = (our local vmprof_tl_stack).next
+        mc.LDR_ri(r.ip1.value, r.sp.value, 4 * WORD)
+        # save in vmprof_tl_stack the value eax
+        offset = cintf.vmprof_tl_stack.getoffset()
+        mc.STR_ri(r.ip1.value, r.ip0.value, offset)
+
     def gen_shadowstack_header(self, gcrootmap):
         # we push two words, like the x86 backend does:
         # the '1' is to benefit from the shadowstack 'is_minor' optimization
diff --git a/rpython/jit/backend/aarch64/runner.py b/rpython/jit/backend/aarch64/runner.py
index ef69fff25d..063b85b483 100644
--- a/rpython/jit/backend/aarch64/runner.py
+++ b/rpython/jit/backend/aarch64/runner.py
@@ -14,6 +14,7 @@ class CPU_ARM64(AbstractLLCPU):
     gen_regs = r.all_regs
     float_regs = VFPRegisterManager.all_regs
     supports_floats = True
+    HAS_CODEMAP = True
 
     from rpython.jit.backend.aarch64.arch import JITFRAME_FIXED_SIZE
 
@@ -29,6 +30,8 @@ class CPU_ARM64(AbstractLLCPU):
 
     def setup_once(self):
         self.assembler.setup_once()
+        if self.HAS_CODEMAP:
+            self.codemap.setup()
 
     def compile_bridge(self, faildescr, inputargs, operations,
                        original_loop_token, log=True, logger=None):
diff --git a/rpython/jit/backend/aarch64/test/test_rvmprof.py b/rpython/jit/backend/aarch64/test/test_rvmprof.py
new file mode 100644
index 0000000000..e2e50e77aa
--- /dev/null
+++ b/rpython/jit/backend/aarch64/test/test_rvmprof.py
@@ -0,0 +1,6 @@
+from rpython.jit.backend.test.test_rvmprof import BaseRVMProfTest
+from rpython.jit.backend.aarch64.test.test_basic import JitAarch64Mixin
+
+
+class TestRVMProfCall(JitAarch64Mixin, BaseRVMProfTest):
+    pass
diff --git a/rpython/jit/backend/aarch64/test/test_zrpy_vmprof.py b/rpython/jit/backend/aarch64/test/test_zrpy_vmprof.py
new file mode 100644
index 0000000000..68b5d6ba03
--- /dev/null
+++ b/rpython/jit/backend/aarch64/test/test_zrpy_vmprof.py
@@ -0,0 +1,7 @@
+
+from rpython.jit.backend.llsupport.test.zrpy_vmprof_test import CompiledVmprofTest
+
+class TestZVMprof(CompiledVmprofTest):
+    
+    gcrootfinder = "shadowstack"
+    gc = "incminimark"
\ No newline at end of file
diff --git a/rpython/rlib/rvmprof/cintf.py b/rpython/rlib/rvmprof/cintf.py
index b8e9492789..37abe6c2d1 100644
--- a/rpython/rlib/rvmprof/cintf.py
+++ b/rpython/rlib/rvmprof/cintf.py
@@ -5,6 +5,7 @@ import shutil
 from rpython.tool.udir import udir
 from rpython.tool.version import rpythonroot
 from rpython.rtyper.lltypesystem import lltype, llmemory, rffi
+from rpython.rtyper.lltypesystem.lloperation import llop
 from rpython.translator.tool.cbuild import ExternalCompilationInfo
 from rpython.rtyper.tool import rffi_platform as platform
 from rpython.rlib import rthread, jit
@@ -19,9 +20,10 @@ class VMProfPlatformUnsupported(Exception):
 IS_SUPPORTED = False
 if sys.platform in ('darwin', 'linux', 'linux2') or sys.platform.startswith('freebsd'):
     try:
-        IS_SUPPORTED = detect_cpu.autodetect().startswith('x86')
+        proc = detect_cpu.autodetect()
+        IS_SUPPORTED = proc.startswith('x86') or proc == 'aarch64'
     except detect_cpu.ProcessorAutodetectError:
-        pass
+        print("PROCESSOR NOT DETECTED, SKIPPING VMPROF")
 
 ROOT = py.path.local(rpythonroot).join('rpython', 'rlib', 'rvmprof')
 SRC = ROOT.join('src')
diff --git a/rpython/rlib/rvmprof/test/__init__.py b/rpython/rlib/rvmprof/test/__init__.py
index 10d19aeb3d..846fe5c449 100644
--- a/rpython/rlib/rvmprof/test/__init__.py
+++ b/rpython/rlib/rvmprof/test/__init__.py
@@ -1,5 +1,5 @@
 import pytest
 import platform
 
-if not platform.machine().startswith('x86'):
+if not (platform.machine().startswith('x86') or platform.machine() == 'aarch64'):
     pytest.skip()
diff --git a/rpython/rlib/rvmprof/test/test_file.py b/rpython/rlib/rvmprof/test/test_file.py
index e13ed978cc..0a14c6dc4b 100644
--- a/rpython/rlib/rvmprof/test/test_file.py
+++ b/rpython/rlib/rvmprof/test/test_file.py
@@ -15,7 +15,10 @@ def get_list_of_files(shared):
     files.remove(shared.join('libbacktrace', 'config-x86_32.h'))
     files.remove(shared.join('libbacktrace', 'config-x86_64.h'))
     files.remove(shared.join('libbacktrace', 'gstdint.h'))
-    files.remove(shared.join('libbacktrace', 'config.h'))
+    try:
+        files.remove(shared.join('libbacktrace', 'config.h'))
+    except ValueError:
+        pass # might not be there
     return files
 
 def test_same_file():