Integrate TraCR as the runtime/kernel profiler for Simpler

.github/workflows/ci.yml

@@ -21,6 +21,7 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v5
         with:
+          submodules: recursive
           fetch-depth: 0
 
       - name: Install build and lint tools
@@ -71,6 +72,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Set up C++ compiler (Linux)
         if: runner.os == 'Linux'
@@ -113,6 +116,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v6
@@ -175,6 +180,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa
@@ -264,6 +271,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa
@@ -354,6 +363,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Set up C++ compiler
         run: |
@@ -454,6 +465,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa
@@ -546,6 +559,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa
@@ -622,6 +637,7 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v5
         with:
+          submodules: recursive
           fetch-depth: 0
       - name: Check file changes
         id: check
@@ -668,6 +684,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa
@@ -743,6 +761,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa

.github/workflows/sanitizers.yml

@@ -38,6 +38,8 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          submodules: recursive
 
       - name: Load pinned pto-isa commit
         uses: ./.github/actions/read-pto-isa

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "tools/tracr"]
+	path = tools/tracr
+	url = https://github.com/huawei-csl/TracR.git

examples/a2a3/tensormap_and_ringbuffer/paged_attention/kernels/orchestration/paged_attention_orch.cpp

@@ -25,6 +25,9 @@
 #include <cstdint>
 #include <cstring>
 
+#include <tracr/tracr.hpp>
+#include <tracr_simpler_markers.hpp>
+
 #include "pto_orchestration_api.h"
 
 #define FUNC_QK_MATMUL 0
@@ -89,6 +92,8 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta
 
     CYCLE_COUNT_START();
 
+    INSTRUMENTATION_MARK_SET(g_TraCR_thread_idx, Read_Dimensions, 0);
+
     // Read dimensions from tensor metadata
     uint64_t batch = orch_args.tensor(0).ref().shapes[0];
     uint64_t num_heads = orch_args.tensor(0).ref().shapes[1];
@@ -107,6 +112,8 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta
 
     LOG_INFO_V9(">>>>>> batch = %" PRIu64, batch);
 
+    INSTRUMENTATION_MARK_SET(g_TraCR_thread_idx, Reshape_Kernels, 0);
+
     // Reshape tensors for kernel consumption (2D flattened)
     void *query_ptr = orch_args.tensor(0).ref().data_as<void>();
     void *kc_ptr = orch_args.tensor(1).ref().data_as<void>();
@@ -136,6 +143,8 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta
     Tensor context_lens =
         make_tensor_external(orch_args.tensor(4).ref().data_as<void>(), cl_shapes, 1, DataType::INT32, false);
 
+    INSTRUMENTATION_MARK_SET(g_TraCR_thread_idx, Pre_Loop_Info, uint32_t(batch));
+
     // Create infos are loop-invariant — shapes depend only on q_tile/head_dim/block_size
     uint32_t tile2d_shapes[2] = {static_cast<uint32_t>(q_tile), static_cast<uint32_t>(head_dim)};
     uint32_t scalar_shapes[1] = {static_cast<uint32_t>(q_tile)};
@@ -148,11 +157,16 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta
     PROF_INC(prof_make_count, 4);
     CYCLE_COUNT_LAP(prof_make_tensor);
 
+    LOG_INFO_V0(
+        "Thread %d: Orch PTO2_SCOPE loop: #batch=%" PRIu64 ", q_loop=%" PRIu64, g_TraCR_thread_idx, batch, q_loop
+    );
+
     for (uint64_t b_idx = 0; b_idx < batch; b_idx++) {
         uint32_t cl_idx[1] = {static_cast<uint32_t>(b_idx)};
         uint64_t cur_seq = static_cast<uint64_t>(get_tensor_data<int32_t>(context_lens, 1, cl_idx));
         uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size;
         for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) {
+            INSTRUMENTATION_MARK_SET(g_TraCR_thread_idx, PTO2_SCOPE_, uint32_t(b_idx + batch * q_idx));
             PTO2_SCOPE() {
                 CYCLE_COUNT_LAP(prof_scope);
                 uint64_t cur_offset = b_idx * q_head_num + q_idx * q_tile;
@@ -250,7 +264,6 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta
             CYCLE_COUNT_LAP(prof_scope);
         }
     }
-
 #ifdef ENABLE_PROFILING
     uint64_t total = prof_param_extract + prof_ext_tensor + prof_make_tensor + prof_tensor_view + prof_param_setup +
                      prof_submit_task + prof_scope;

simpler_setup/kernel_compiler.py

@@ -143,7 +143,17 @@ def get_orchestration_include_dirs(self, runtime_name: str) -> list[str]:
         runtime_dir = str(self.project_root / "src" / arch / "runtime" / runtime_name / "runtime")
         runtime_common_dir = str(self.project_root / "src" / arch / "runtime" / runtime_name / "common")
         common_dir = str(self.project_root / "src" / "common" / "task_interface")
-        return [runtime_dir, runtime_common_dir, common_dir] + self.get_platform_include_dirs()
+        tracr_dir1 = str(self.project_root / "tools")
+        tracr_dir2 = str(self.project_root / "tools" / "tracr" / "include")
+        tracr_dir3 = str(self.project_root / "tools" / "tracr" / "extern")
+        return [
+            runtime_dir,
+            runtime_common_dir,
+            common_dir,
+            tracr_dir1,
+            tracr_dir2,
+            tracr_dir3,
+        ] + self.get_platform_include_dirs()
 
     def get_incore_include_dirs(self) -> list[str]:
         """
@@ -487,6 +497,15 @@ def _compile_orchestration_shared_lib(
         if sys.platform != "darwin":
             cmd.append("-Wl,--build-id=sha1")
 
+        if os.getenv("BUILD_TRACR", "OFF") == "ON":
+            cmd.extend(
+                [
+                    "-DENABLE_TRACR",
+                    "-DTRACR_DISABLE_FLUSH",
+                    "-DUSE_HW_COUNTER",
+                ]
+            )
+
         if extra_sources:
             for src in extra_sources:
                 src = os.path.abspath(src)

src/a2a3/platform/onboard/aicpu/CMakeLists.txt

@@ -63,6 +63,19 @@ message(STATUS "AICpu kernel: ${NUM_AICPU_KERNEL_SOURCES} source files")
 message(VERBOSE "AICpu kernel sources: ${AICPU_KERNEL_SOURCES}")
 add_library(aicpu_kernel SHARED ${AICPU_KERNEL_SOURCES})
 
+# TraCR
+include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/tracr.cmake)
+tracr_enable(aicpu_kernel)
+
+# Option to make the Orchestrator run independently (i.e. finalize before letting the schedulers run)
+option(INDEP_ORCH "Run Orchestrator independent from the Schedulers" OFF)
+if(DEFINED ENV{INDEP_ORCH})
+    set(INDEP_ORCH $ENV{INDEP_ORCH})
+endif()
+if(INDEP_ORCH)
+    target_compile_definitions(aicpu_kernel PRIVATE INDEP_ORCH)
+endif()
+
 option(WERROR "Treat compiler warnings as errors" ON)
 if(DEFINED ENV{SIMPLER_DISABLE_WARNINGS_AS_ERRORS})
     set(WERROR OFF)

src/a2a3/platform/onboard/host/CMakeLists.txt

@@ -91,6 +91,19 @@ message(STATUS "Host runtime: ${NUM_HOST_RUNTIME_SOURCES} source files")
 message(VERBOSE "Host runtime sources: ${HOST_RUNTIME_SOURCES}")
 add_library(host_runtime SHARED ${HOST_RUNTIME_SOURCES})
 
+# TraCR
+include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/tracr.cmake)
+tracr_enable(host_runtime)
+
+# Optional" to make the Orchestrator run independently (i.e. finalize before letting the schedulers run)
+option(INDEP_ORCH "Run Orchestrator independent from the Schedulers" OFF)
+if(DEFINED ENV{INDEP_ORCH})
+    set(INDEP_ORCH $ENV{INDEP_ORCH})
+endif()
+if(INDEP_ORCH)
+    target_compile_definitions(host_runtime PRIVATE INDEP_ORCH)
+endif()
+
 # C++ standard (applied only to C++ files)
 set_target_properties(host_runtime PROPERTIES
     CXX_STANDARD 17

src/a2a3/platform/onboard/host/device_runner.cpp

@@ -28,6 +28,9 @@
 #include <iostream>
 #include <string>
 #include <vector>
+
+#include <tracr_simpler_api.hpp>
+
 #include "acl/acl.h"
 
 // Include HAL constants from CANN (header only, library loaded dynamically)
@@ -264,6 +267,15 @@ int DeviceRunner::run(Runtime &runtime, int block_dim, int launch_aicpu_num) {
 
     if (prepare_runtime_for_launch(runtime, block_dim, launch_aicpu_num) != 0) return -1;
 
+    // Initialize TraCR memory on the device
+#ifdef ENABLE_TRACR
+    // LOG_INFO_V9("[TraCR] thread[%d] DevAllocTraCR device_id_=%d", sched_getcpu(), device_id_);
+    rc = DevAllocTraCR(this, runtime);
+    if (rc != 0) {
+        LOG_ERROR("DevAllocTraCR failed rc=%d", rc);
+        return rc;
+    }
+#endif
     // a2a3 onboard now uses the same host-computed, device-filtered affinity
     // shape as a5. Host probes the AICPU user pool once, chooses the active
     // cpu_ids deterministically, writes them into Runtime, and the AICPU-side
@@ -471,6 +483,15 @@ int DeviceRunner::run(Runtime &runtime, int block_dim, int launch_aicpu_num) {
 
     read_device_wall_ns();
 
+    // Download and Free TraCR memory from Device and store in memory (~/ascend/)
+#ifdef ENABLE_TRACR
+    rc = StoreTracrData(this, runtime);
+    if (rc != 0) {
+        LOG_ERROR("FreeTraCR failed: %d", rc);
+        return -1;
+    }
+#endif
+
     // Tear down collectors. stop() joins mgmt then collector in the only safe
     // order (mgmt's final-drain pass into L2 has poll as its consumer).
     teardown_shared_collectors_after_run();

src/a2a3/platform/sim/aicpu/CMakeLists.txt

@@ -70,6 +70,27 @@ endif()
 # Create shared library (host-compatible for dlopen)
 add_library(aicpu_kernel SHARED ${AICPU_SOURCES})
 
+# TraCR
+include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/tracr.cmake)
+tracr_enable(aicpu_kernel)
+
+# TODO: move this somewhere such that EVERY platform launches this once. Placing this here is hacky...
+# Only build the host-side trace post-processor when TraCR is enabled: it is an
+# offline analysis tool, and it pulls in Linux-only APIs (sched_getcpu) that do
+# not compile on the macOS packaging build.
+if(BUILD_TRACR)
+    include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/tracr_postprocessing_script.cmake)
+endif()
+
+# Optional: to make the Orchestrator run independently (i.e. finalize before letting the schedulers run)
+option(INDEP_ORCH "Run Orchestrator independent from the Schedulers" OFF)
+if(DEFINED ENV{INDEP_ORCH})
+    set(INDEP_ORCH $ENV{INDEP_ORCH})
+endif()
+if(INDEP_ORCH)
+    target_compile_definitions(aicpu_kernel PRIVATE INDEP_ORCH)
+endif()
+
 option(WERROR "Treat compiler warnings as errors" ON)
 if(DEFINED ENV{SIMPLER_DISABLE_WARNINGS_AS_ERRORS})
     set(WERROR OFF)

src/a2a3/platform/sim/host/CMakeLists.txt

@@ -75,6 +75,19 @@ endif()
 # Create shared library
 add_library(host_runtime SHARED ${HOST_RUNTIME_SOURCES})
 
+# TraCR
+include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/tracr.cmake)
+tracr_enable(host_runtime)
+
+# Optional: to make the Orchestrator run independently (i.e. finalize before letting the schedulers run)
+option(INDEP_ORCH "Run Orchestrator independent from the Schedulers" OFF)
+if(DEFINED ENV{INDEP_ORCH})
+    set(INDEP_ORCH $ENV{INDEP_ORCH})
+endif()
+if(INDEP_ORCH)
+    target_compile_definitions(host_runtime PRIVATE INDEP_ORCH)
+endif()
+
 # C++ standard (applied only to C++ files)
 set_target_properties(host_runtime PROPERTIES
     CXX_STANDARD 17

src/a2a3/platform/sim/host/device_runner.cpp

@@ -29,6 +29,8 @@
 #include <string>
 #include <vector>
 
+#include <tracr_simpler_api.hpp>
+
 #include "aicpu/device_phase_aicpu.h"
 #include "aicpu/platform_aicpu_affinity.h"
 #include "callable_protocol.h"
@@ -253,6 +255,15 @@ int DeviceRunner::run(Runtime &runtime, int block_dim, int launch_aicpu_num) {
     worker_count_ = num_aicore;
     runtime.set_aicpu_thread_num(launch_aicpu_num);
 
+    // Initialize TraCR memory on the device
+#ifdef ENABLE_TRACR
+    rc = DevAllocTraCR(this, runtime);
+    if (rc != 0) {
+        LOG_ERROR("DevAllocTraCR failed rc=%d", rc);
+        return rc;
+    }
+#endif
+
     int num_aic = block_dim;
     uint32_t enable_profiling_flag = PROFILING_FLAG_NONE;
     if (enable_dump_tensor_) {
@@ -571,6 +582,15 @@ int DeviceRunner::run(Runtime &runtime, int block_dim, int launch_aicpu_num) {
         return runtime_rc;
     }
 
+    // Download and Free TraCR memory from Device and store in memory (~/ascend/)
+#ifdef ENABLE_TRACR
+    rc = StoreTracrData(this, runtime);
+    if (rc != 0) {
+        LOG_ERROR("FreeTraCR failed: %d", rc);
+        return -1;
+    }
+#endif
+
     // Tear down collectors. stop() joins mgmt then collector in the only safe
     // order (mgmt's final-drain pass into L2 has poll as its consumer).
     if (enable_l2_swimlane_) {

src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h

@@ -20,7 +20,12 @@
 
 // Tensor dump uses these defaults to size its selective mask table so task-id
 // ring/slot lookup stays aligned with PTO2 task id layout.
+#ifdef INDEP_ORCH
+#define PTO2_TASK_WINDOW_SIZE 65536  // Default per-ring task window size (power of 2)
+#else
 #define PTO2_TASK_WINDOW_SIZE 16384  // Default per-ring task window size (power of 2)
-#define PTO2_MAX_RING_DEPTH 4        // Number of task-id ring layers
+#endif
+
+#define PTO2_MAX_RING_DEPTH 4  // Number of task-id ring layers
 
 #endif  // SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_PTO_RUNTIME2_TYPES_H_