Make asan instrumentation in mj_(mark|free)Stack cheaper.

Symbolization and string comparison results are now cached.

PiperOrigin-RevId: 564697198
Change-Id: I0447f283a90143e2f67ecd90d895a36a078b3981

src/engine/CMakeLists.txt

@@ -29,7 +29,7 @@ set(MUJOCO_ENGINE_SRCS
     engine_core_constraint.h
     engine_core_smooth.c
     engine_core_smooth.h
-    engine_crossplatform.c
+    engine_crossplatform.cc
     engine_crossplatform.h
     engine_derivative.c
     engine_derivative.h

src/engine/engine_crossplatform.c

@@ -1,41 +0,0 @@
-// Copyright 2022 DeepMind Technologies Limited
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-void _mj_crossplatform_void(void) {}  // ISO C does not permit empty translation units
-
-#if defined(__APPLE__) && defined(__AVX__)
-
-#include <stdio.h>
-#include <string.h>
-#include <sys/sysctl.h>
-
-__attribute__((weak, visibility("default"))) void _mj_rosettaError(const char* msg) {
-  fprintf(stderr, "%s\n", msg);
-  __asm__ __volatile__ ("ud2");  // raises SIGILL but leave this function at the top of the stack
-}
-
-__attribute__((constructor(10000), target("no-avx"))) static void _mj_checkRosetta(void) {
-  int is_translated = 0;
-  {
-    size_t len = sizeof(is_translated);
-    if (sysctlbyname("sysctl.proc_translated", &is_translated, &len, NULL, 0)) {
-      is_translated = 0;
-    }
-  }
-  if (is_translated) {
-    _mj_rosettaError("MuJoCo cannot be run under Rosetta 2 on an Apple Silicon machine.");
-  }
-}
-
-#endif  // defined(__APPLE__) && defined(__AVX__)

src/engine/engine_crossplatform.cc

@@ -0,0 +1,99 @@
+// Copyright 2022 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_crossplatform.h"  // IWYU pragma: keep
+
+#if defined(__APPLE__) && defined(__AVX__)
+#include <sys/sysctl.h>
+
+#include <cstdio>
+#include <cstring>
+
+namespace {
+__attribute__((weak, visibility("default")))
+extern "C" void _mj_rosettaError(const char* msg) {
+  fprintf(stderr, "%s\n", msg);
+  __asm__ __volatile__ ("ud2");  // raises SIGILL but leave this function at the top of the stack
+}
+
+__attribute__((constructor(10000), target("no-avx")))
+void CheckRosetta() {
+  int is_translated = 0;
+  {
+    size_t len = sizeof(is_translated);
+    if (sysctlbyname("sysctl.proc_translated", &is_translated, &len, NULL, 0)) {
+      is_translated = 0;
+    }
+  }
+  if (is_translated) {
+    _mj_rosettaError("MuJoCo cannot be run under Rosetta 2 on an Apple Silicon machine.");
+  }
+}
+}  // namespace
+#endif  // defined(__APPLE__) && defined(__AVX__)
+
+#ifdef ADDRESS_SANITIZER
+#include <sanitizer/common_interface_defs.h>
+
+#include <array>
+#include <map>
+#include <mutex>
+#include <shared_mutex>
+#include <string>
+#include <string_view>
+#include <unordered_map>
+#include <utility>
+
+namespace {
+std::string_view SymbolizeCached(void* pc) {
+  static auto* mu = new std::shared_mutex;
+  static auto* pc_to_func_name_map = new std::unordered_map<void*, std::string>;
+
+  {
+    std::shared_lock lock(*mu);
+    auto it = pc_to_func_name_map->find(pc);
+    if (it != pc_to_func_name_map->end()) {
+      return it->second;
+    }
+  }
+
+  std::array<char, 256> buf;
+  __sanitizer_symbolize_pc(pc, "%f", buf.data(), buf.size());
+  {
+    std::unique_lock lock(*mu);
+    return pc_to_func_name_map->emplace(pc, buf.data()).first->second;
+  }
+}
+}  // namespace
+
+int _mj_comparePcFuncName(void* pc1, void* pc2) {
+  static auto* mu = new std::shared_mutex;
+  static auto* same_func_map = new std::map<std::pair<void*, void*>, bool>;
+
+  auto pc_pair = std::make_pair(pc1, pc2);
+  {
+    std::shared_lock lock(*mu);
+    auto it = same_func_map->find(pc_pair);
+    if (it != same_func_map->end()) {
+      return it->second;
+    }
+  }
+
+  bool is_same = (SymbolizeCached(pc1) == SymbolizeCached(pc2));
+  {
+    std::unique_lock lock(*mu);
+    return same_func_map->emplace(pc_pair, is_same).first->second;
+  }
+}
+#endif  // ADDRESS_SANITIZER

src/engine/engine_crossplatform.h

@@ -75,4 +75,16 @@
   #define mjUNLIKELY(x) (x)
 #endif
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifdef ADDRESS_SANITIZER
+int _mj_comparePcFuncName(void* pc1, void* pc2);
+#endif
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_CROSSPLATFORM_H_

src/engine/engine_io.c

@@ -1351,29 +1351,17 @@ void mj_freeStack(mjData* d) {
 
   mjStackFrame* s = (mjStackFrame*) ((char*)d->arena + d->narena - d->pbase);
 #ifdef ADDRESS_SANITIZER
-  #define mjSYMBOLIZELEN 256
-
-  // symbolize s->pc to get the function name of most recent caller to mj_markStack
-  char markstack_func[mjSYMBOLIZELEN];
-  __sanitizer_symbolize_pc(s->pc, "%f", markstack_func, mjSYMBOLIZELEN);
-  markstack_func[mjSYMBOLIZELEN - 1] = '\0';
-
-  // symbolize current program counter to get the function name of caller to this function
-  char freestack_func[mjSYMBOLIZELEN];
-  __sanitizer_symbolize_pc(__sanitizer_return_address(), "%f", freestack_func, mjSYMBOLIZELEN);
-  freestack_func[mjSYMBOLIZELEN - 1] = '\0';
-
   // raise an error if caller function name doesn't match the most recent caller of mj_markStack
-  if (strncmp(markstack_func, freestack_func, mjSYMBOLIZELEN)) {
+  if (!_mj_comparePcFuncName(s->pc, __sanitizer_return_address())) {
+    #define mjSYMBOLIZELEN 256
     char dbginfo[mjSYMBOLIZELEN];
     __sanitizer_symbolize_pc(
         s->pc, "mj_markStack %F at %S has no corresponding mj_freeStack",
         dbginfo, sizeof(dbginfo));
     dbginfo[mjSYMBOLIZELEN - 1] = '\0';
     mjERROR("%s", dbginfo);
+    #undef mjSYMBOLIZELEN
   }
-
-  #undef mjSYMBOLIZELEN
 #endif
 
   // restore pbase and pstack