Use tablegen Python bindings for trace events instead of separate generation

include/aie/Dialect/AIE/IR/CMakeLists.txt

@@ -60,21 +60,29 @@ set(GENERATED_EVENTS_TD_FILES
   "${CMAKE_BINARY_DIR}/include/AIEEventsAIE2P.td.inc"
 )
 
+set(GENERATED_EVENTS_JSON "${CMAKE_BINARY_DIR}/lib/regdb/events_database.json")
+
 add_custom_command(
-  OUTPUT ${GENERATED_EVENTS_TD_FILES}
+  OUTPUT ${GENERATED_EVENTS_TD_FILES} ${GENERATED_EVENTS_JSON}
   COMMAND ${Python3_EXECUTABLE}
     ${AIE_SOURCE_DIR}/utils/generate_events_tablegen.py
     -i ${EVENTS_HEADER_FILES}
     -o ${CMAKE_BINARY_DIR}/include
+  COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_BINARY_DIR}/lib/regdb
+  COMMAND ${Python3_EXECUTABLE}
+    ${AIE_SOURCE_DIR}/utils/generate_events_json.py
+    -i ${EVENTS_HEADER_FILES}
+    -o ${CMAKE_BINARY_DIR}/lib/regdb
   DEPENDS
     ${AIE_SOURCE_DIR}/utils/generate_events_tablegen.py
+    ${AIE_SOURCE_DIR}/utils/generate_events_json.py
     ${EVENTS_HEADER_FILES}
-  COMMENT "Generating AIE event enum tablegen files from aie-rt headers"
+  COMMENT "Generating AIE event tablegen and JSON database from aie-rt headers"
   VERBATIM
 )
 
 add_custom_target(GenerateAIEEventsTD
-  DEPENDS ${GENERATED_EVENTS_TD_FILES}
+  DEPENDS ${GENERATED_EVENTS_TD_FILES} ${GENERATED_EVENTS_JSON}
 )
 
 # Add AIE attributes

lib/Dialect/AIE/Util/CMakeLists.txt

@@ -36,8 +36,9 @@ add_mlir_dialect_library(MLIRAIEUtil
   MLIRSupport
 )
 
-add_dependencies(MLIRAIEUtil MLIRAIERegDBResources GenerateTraceEventsEnum)
+add_dependencies(MLIRAIEUtil MLIRAIERegDBResources GenerateAIEEventsTD)
 
 install(FILES
   ${AIE_BINARY_DIR}/lib/regdb/aie_registers_aie2.json
+  ${AIE_BINARY_DIR}/lib/regdb/events_database.json
   DESTINATION lib${LLVM_LIBDIR_SUFFIX}/regdb)

python/CMakeLists.txt

@@ -387,85 +387,10 @@ else ()
     COMPONENT aie-python
   )
 
-  ################################################################################
-  # Generated trace events enum files and copy to Python package directories
-  ################################################################################
-
-  # Input header files from aie-rt
-  set(AIE_RT_EVENTS_DIR "${AIE_SOURCE_DIR}/third_party/aie-rt/driver/src/events")
-  set(EVENTS_HEADER_FILES
-    "${AIE_RT_EVENTS_DIR}/xaie_events_aie.h"
-    "${AIE_RT_EVENTS_DIR}/xaie_events_aieml.h"
-    "${AIE_RT_EVENTS_DIR}/xaie_events_aie2p.h"
-  )
-
-  # Output directory for generated files in build directory
-  set(GENERATED_EVENTS_DIR "${CMAKE_BINARY_DIR}/lib/regdb")
-
-  # Generated Python enum files in the build directory
-  set(GENERATED_EVENTS_ENUM_FILES
-    "${GENERATED_EVENTS_DIR}/aie.py"
-    "${GENERATED_EVENTS_DIR}/aie2.py"
-    "${GENERATED_EVENTS_DIR}/aie2p.py"
-  )
-
-  # Generated JSON database in the build directory
-  set(GENERATED_EVENTS_JSON_FILENAME "events_database.json")
-  set(GENERATED_EVENTS_JSON "${GENERATED_EVENTS_DIR}/${GENERATED_EVENTS_JSON_FILENAME}")
-
-  # Add custom command to generate the event enum files
-  add_custom_command(
-    OUTPUT ${GENERATED_EVENTS_ENUM_FILES} ${GENERATED_EVENTS_JSON}
-    COMMAND ${CMAKE_COMMAND} -E make_directory ${GENERATED_EVENTS_DIR}
-    COMMAND ${Python3_EXECUTABLE} 
-      ${AIE_SOURCE_DIR}/utils/generate_events_enum.py
-      -i ${EVENTS_HEADER_FILES}
-      -o ${GENERATED_EVENTS_DIR}
-      --json ${GENERATED_EVENTS_JSON_FILENAME}
-      --python-prefix ""
-    DEPENDS 
-      ${AIE_SOURCE_DIR}/utils/generate_events_enum.py
-      ${EVENTS_HEADER_FILES}
-    COMMENT "Generating trace events enum files from aie-rt headers"
-    VERBATIM
-  )
-
-  # Install the generated json to the lib dir
-  install(
-    FILES 
-      ${GENERATED_EVENTS_JSON}
-    DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/regdb
-    COMPONENT aie-python
-  )
-
-  # Create a custom target for the generated files
-  add_custom_target(GenerateTraceEventsEnum ALL
-    DEPENDS ${GENERATED_EVENTS_ENUM_FILES} ${GENERATED_EVENTS_JSON}
-  )
-
-  # Ensure AIEPythonModules depends on GenerateTraceEventsEnum
-  add_dependencies(AIEPythonModules GenerateTraceEventsEnum)
-
-  # Copy generated event enum files to Python package build directory
-  foreach(_file ${GENERATED_EVENTS_ENUM_FILES})
-    get_filename_component(_filename ${_file} NAME)
-    add_custom_command(
-      TARGET AIEPythonModules POST_BUILD
-      COMMAND ${CMAKE_COMMAND} -E make_directory ${AIE_PYTHON_PACKAGES_DIR}/aie/utils/trace/events
-      COMMAND ${CMAKE_COMMAND} -E copy_if_different
-        ${_file}
-        ${AIE_PYTHON_PACKAGES_DIR}/aie/utils/trace/events/${_filename}
-      COMMENT "Copying ${_filename} to build Python package"
-    )
-  endforeach()
-
-  # Install generated event enum files
-  install(
-    FILES 
-      ${GENERATED_EVENTS_ENUM_FILES}
-    DESTINATION ${AIE_PYTHON_INSTALL_DIR}/aie/utils/trace/events
-    COMPONENT aie-python
-  )
+  # Note: events_database.json is generated at build time by the
+  # GenerateAIEEventsTD target (include/aie/Dialect/AIE/IR/CMakeLists.txt).
+  # Python trace event enums are sourced from the TableGen-generated
+  # _aie_enum_gen.py bindings (no separate generation step needed).
 
 
 endif ()

python/utils/trace/events/__init__.py

@@ -6,27 +6,41 @@
 #
 """Trace events enumerations for AIE architectures.
 
-Available modules:
-- aie: AIE1 architecture events
-- aie2: AIE2/AIEML architecture events
-- aie2p: AIE2P architecture events
+Event enums are sourced from the TableGen-generated Python bindings
+(aie.dialects._aie_enum_gen), which are produced from the same aie-rt
+headers that define the hardware event numbers.
+
+Architecture-specific enums (CoreEventAIE2, etc.) are re-exported here
+under architecture-agnostic names (CoreEvent, etc.) for convenience.
+Use get_events_for_device() to select the correct architecture.
 """
+
 from enum import IntEnum
+from types import SimpleNamespace
 import typing
 
-from . import aie
-from . import aie2
-from . import aie2p
-
-from .aie2 import (
-    CoreEvent,
-    MemEvent,
-    ShimTileEvent,
-    MemTileEvent,
+from aie.dialects._aie_enum_gen import (
+    CoreEventAIE,
+    MemEventAIE,
+    ShimTileEventAIE,
+    CoreEventAIE2,
+    MemEventAIE2,
+    ShimTileEventAIE2,
+    MemTileEventAIE2,
+    CoreEventAIE2P,
+    MemEventAIE2P,
+    ShimTileEventAIE2P,
+    MemTileEventAIE2P,
 )
 
 from aie.dialects.aie import WireBundle, DMAChannelDir
 
+# Default to AIE2 for backwards compatibility
+CoreEvent = CoreEventAIE2
+MemEvent = MemEventAIE2
+ShimTileEvent = ShimTileEventAIE2
+MemTileEvent = MemTileEventAIE2
+
 
 # We use the packet type field in the packet header to help differentiate the tile
 # that the packet came from. Since packet types don't inherently have meaning, we
@@ -44,11 +58,26 @@ class PacketType(IntEnum):
 
 def get_events_for_device(device: str):
     if "xcvc1902" in device:
-        return aie
+        return SimpleNamespace(
+            CoreEvent=CoreEventAIE,
+            MemEvent=MemEventAIE,
+            ShimTileEvent=ShimTileEventAIE,
+            MemTileEvent=None,  # AIE1 has no mem tiles
+        )
     elif "npu2p" in device:
-        return aie2p
+        return SimpleNamespace(
+            CoreEvent=CoreEventAIE2P,
+            MemEvent=MemEventAIE2P,
+            ShimTileEvent=ShimTileEventAIE2P,
+            MemTileEvent=MemTileEventAIE2P,
+        )
     else:
-        return aie2
+        return SimpleNamespace(
+            CoreEvent=CoreEventAIE2,
+            MemEvent=MemEventAIE2,
+            ShimTileEvent=ShimTileEventAIE2,
+            MemTileEvent=MemTileEventAIE2,
+        )
 
 
 def _get_port_events(enum_class):
@@ -68,9 +97,6 @@ class GenericEvent:
     def __init__(
         self, code: typing.Union[CoreEvent, MemEvent, ShimTileEvent, MemTileEvent]
     ):
-        # For backwards compatibility, allow integer as event
-        if isinstance(code, int):
-            code = CoreEvent(code)
         self.code: typing.Union[CoreEvent, MemEvent, ShimTileEvent, MemTileEvent] = code
 
     def get_register_writes(self):
@@ -118,9 +144,6 @@ def __init__(
         enum_class=None,
         valid_codes=None,
     ):
-        # For backwards compatibility, allow integer as event
-        if isinstance(code, int) and enum_class:
-            code = enum_class(code)
         if valid_codes:
             assert code in valid_codes
 

python/utils/trace/parse.py

@@ -551,6 +551,9 @@ def lookup_event_name_by_type(trace_type, code, events_module):
     else:
         return "Unknown"
 
+    if enum_class is None:
+        return "Unknown"
+
     try:
         return enum_class(code).name
     except ValueError:

test/python/trace_events.py

@@ -26,16 +26,19 @@
 # CHECK: Edge case: Arch-specific events
 # CHECK: AIE1.MemEvent has WATCHPOINT_0
 # CHECK: AIE2.MemEvent has WATCHPOINT_0
-# CHECK: AIE2 has MemTileEvent, AIE1 does not
-# CHECK: AIE2P has MemTileEvent, AIE1 does not
+# CHECK: AIE1 MemTileEvent is None (no mem tiles in AIE1)
+# CHECK: AIE2 has MemTileEvent
+# CHECK: AIE2P has MemTileEvent
 # CHECK: Edge case: Events not present in arch raise AttributeError
 # CHECK: AIE1.CoreEvent.EDGE_DETECTION_EVENT_0 does not exist
 # CHECK: AIE2.CoreEvent.SRS_SATURATE does not exist
 # CHECK: All tests passed!
 
-import aie.utils.trace.events.aie as aie1
-import aie.utils.trace.events.aie2 as aie2
-import aie.utils.trace.events.aie2p as aie2p
+from aie.utils.trace.events import get_events_for_device
+
+aie1 = get_events_for_device("xcvc1902")
+aie2 = get_events_for_device("npu1")
+aie2p = get_events_for_device("npu2p")
 
 
 def test_arch(name, core_cls, mem_cls, shim_cls, mem_tile_cls=None):
@@ -132,22 +135,18 @@ def test_edge_cases():
     assert hasattr(aie2.MemEvent, "WATCHPOINT_0"), "AIE2 should have WATCHPOINT_0"
     print("AIE2.MemEvent has WATCHPOINT_0")
 
-    # MemTileEvent only exists in AIE2/AIE2P
-    assert hasattr(aie1, "MemTileEvent"), "AIE1 module has MemTileEvent class"
-    # But AIE1's MemTileEvent has no events (just pass)
-    aie1_memtile_events = [e for e in aie1.MemTileEvent]
-    assert len(aie1_memtile_events) == 0, "AIE1 MemTileEvent should be empty"
-    print("AIE2 has MemTileEvent, AIE1 does not")
+    # AIE1 has no mem tiles
+    assert aie1.MemTileEvent is None, "AIE1 should not have MemTileEvent"
+    print("AIE1 MemTileEvent is None (no mem tiles in AIE1)")
 
     # AIE2/AIE2P have actual MemTileEvent events
-    assert hasattr(aie2, "MemTileEvent"), "AIE2 should have MemTileEvent"
     aie2_memtile_events = [e for e in aie2.MemTileEvent]
     assert len(aie2_memtile_events) > 0, "AIE2 MemTileEvent should have events"
+    print("AIE2 has MemTileEvent")
 
-    assert hasattr(aie2p, "MemTileEvent"), "AIE2P should have MemTileEvent"
     aie2p_memtile_events = [e for e in aie2p.MemTileEvent]
     assert len(aie2p_memtile_events) > 0, "AIE2P MemTileEvent should have events"
-    print("AIE2P has MemTileEvent, AIE1 does not")
+    print("AIE2P has MemTileEvent")
 
     print("Edge case: Events not present in arch raise AttributeError")