[mypyc] Optionally log a sampled operation trace to a file

mypyc/__main__.py

@@ -23,8 +23,15 @@
 from setuptools import setup
 from mypyc.build import mypycify
 
-setup(name='mypyc_output',
-      ext_modules=mypycify({}, opt_level="{}", debug_level="{}", strict_dunder_typing={}),
+setup(
+    name='mypyc_output',
+    ext_modules=mypycify(
+        {},
+        opt_level="{}",
+        debug_level="{}",
+        strict_dunder_typing={},
+        log_trace={},
+    ),
 )
 """
 
@@ -39,10 +46,17 @@ def main() -> None:
     opt_level = os.getenv("MYPYC_OPT_LEVEL", "3")
     debug_level = os.getenv("MYPYC_DEBUG_LEVEL", "1")
     strict_dunder_typing = bool(int(os.getenv("MYPYC_STRICT_DUNDER_TYPING", "0")))
+    # If enabled, compiled code writes a sampled log of executed ops (or events) to
+    # mypyc_trace.txt.
+    log_trace = bool(int(os.getenv("MYPYC_LOG_TRACE", "0")))
 
     setup_file = os.path.join(build_dir, "setup.py")
     with open(setup_file, "w") as f:
-        f.write(setup_format.format(sys.argv[1:], opt_level, debug_level, strict_dunder_typing))
+        f.write(
+            setup_format.format(
+                sys.argv[1:], opt_level, debug_level, strict_dunder_typing, log_trace
+            )
+        )
 
     # We don't use run_setup (like we do in the test suite) because it throws
     # away the error code from distutils, and we don't care about the slight

mypyc/build.py

@@ -485,6 +485,7 @@ def mypycify(
     include_runtime_files: bool | None = None,
     strict_dunder_typing: bool = False,
     group_name: str | None = None,
+    log_trace: bool = False,
 ) -> list[Extension]:
     """Main entry point to building using mypyc.
 
@@ -531,6 +532,10 @@ def mypycify(
                     the hash of module names. This is used for the names of the
                     output C files and the shared library. This is only supported
                     if there is a single group. [Experimental]
+        log_trace: If True, compiled code writes a trace log of events in
+                   mypyc_trace.txt (derived from executed operations). This is
+                   useful for performance analysis, such as analyzing which
+                   primitive ops are used the most and on which lines.
     """
 
     # Figure out our configuration
@@ -543,6 +548,7 @@ def mypycify(
         include_runtime_files=include_runtime_files,
         strict_dunder_typing=strict_dunder_typing,
         group_name=group_name,
+        log_trace=log_trace,
     )
 
     # Generate all the actual important C code
@@ -583,6 +589,8 @@ def mypycify(
             # See https://github.com/mypyc/mypyc/issues/956
             "-Wno-cpp",
         ]
+        if log_trace:
+            cflags.append("-DMYPYC_LOG_TRACE")
     elif compiler.compiler_type == "msvc":
         # msvc doesn't have levels, '/O2' is full and '/Od' is disable
         if opt_level == "0":
@@ -607,6 +615,8 @@ def mypycify(
             # that we actually get the compilation speed and memory
             # use wins that multi-file mode is intended for.
             cflags += ["/GL-", "/wd9025"]  # warning about overriding /GL
+        if log_trace:
+            cflags.append("/DMYPYC_LOG_TRACE")
 
     # If configured to (defaults to yes in multi-file mode), copy the
     # runtime library in. Otherwise it just gets #included to save on

mypyc/codegen/emitmodule.py

@@ -65,6 +65,7 @@
 from mypyc.transform.copy_propagation import do_copy_propagation
 from mypyc.transform.exceptions import insert_exception_handling
 from mypyc.transform.flag_elimination import do_flag_elimination
+from mypyc.transform.log_trace import insert_event_trace_logging
 from mypyc.transform.lower import lower_ir
 from mypyc.transform.refcount import insert_ref_count_opcodes
 from mypyc.transform.spill import insert_spills
@@ -253,6 +254,9 @@ def compile_scc_to_ir(
             if fn in env_user_functions:
                 insert_spills(fn, env_user_functions[fn])
 
+            if compiler_options.log_trace:
+                insert_event_trace_logging(fn, compiler_options)
+
             # Switch to lower abstraction level IR.
             lower_ir(fn, compiler_options)
             # Perform optimizations.

mypyc/lib-rt/CPy.h

@@ -928,6 +928,7 @@ PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, PyOb
 PyObject *CPy_GetAIter(PyObject *obj);
 PyObject *CPy_GetANext(PyObject *aiter);
 void CPy_SetTypeAliasTypeComputeFunction(PyObject *alias, PyObject *compute_value);
+void CPyTrace_LogEvent(const char *location, const char *line, const char *op, const char *details);
 
 #ifdef __cplusplus
 }

mypyc/lib-rt/misc_ops.c

@@ -1030,6 +1030,34 @@ PyObject *CPy_GetANext(PyObject *aiter)
     return NULL;
 }
 
+#ifdef MYPYC_LOG_TRACE
+
+// This is only compiled in if trace logging is enabled by user
+
+static int TraceCounter = 0;
+static const int TRACE_EVERY_NTH = 1009;  // Should be a prime number
+#define TRACE_LOG_FILE_NAME "mypyc_trace.txt"
+static FILE *TraceLogFile = NULL;
+
+// Log a tracing event on every Nth call
+void CPyTrace_LogEvent(const char *location, const char *line, const char *op, const char *details) {
+    if (TraceLogFile == NULL) {
+        if ((TraceLogFile = fopen(TRACE_LOG_FILE_NAME, "w")) == NULL) {
+            fprintf(stderr, "error: Could not open trace file %s\n", TRACE_LOG_FILE_NAME);
+            abort();
+        }
+    }
+    if (TraceCounter == 0) {
+        fprintf(TraceLogFile, "%s:%s:%s:%s\n", location, line, op, details);
+    }
+    TraceCounter++;
+    if (TraceCounter == TRACE_EVERY_NTH) {
+        TraceCounter = 0;
+    }
+}
+
+#endif
+
 #ifdef CPY_3_12_FEATURES
 
 // Copied from Python 3.12.3, since this struct is internal to CPython. It defines

mypyc/options.py

@@ -16,6 +16,7 @@ def __init__(
         python_version: tuple[int, int] | None = None,
         strict_dunder_typing: bool = False,
         group_name: str | None = None,
+        log_trace: bool = False,
     ) -> None:
         self.strip_asserts = strip_asserts
         self.multi_file = multi_file
@@ -45,3 +46,7 @@ def __init__(
         # library is generated (with shims). This can be used to make the output
         # file names more predictable.
         self.group_name = group_name
+        # If enabled, write a trace log of events based on executed operations to
+        # mypyc_trace.txt when compiled module is executed. This is useful for
+        # performance analysis.
+        self.log_trace = log_trace

mypyc/primitives/misc_ops.py

@@ -9,6 +9,7 @@
     c_int_rprimitive,
     c_pointer_rprimitive,
     c_pyssize_t_rprimitive,
+    cstring_rprimitive,
     dict_rprimitive,
     int_rprimitive,
     object_pointer_rprimitive,
@@ -291,3 +292,13 @@
     return_type=void_rtype,
     error_kind=ERR_NEVER,
 )
+
+# Log an event to a trace log, which is written to a file during execution.
+log_trace_event = custom_primitive_op(
+    name="log_trace_event",
+    c_function_name="CPyTrace_LogEvent",
+    # (fullname of function/location, line number, operation name, operation details)
+    arg_types=[cstring_rprimitive, cstring_rprimitive, cstring_rprimitive, cstring_rprimitive],
+    return_type=void_rtype,
+    error_kind=ERR_NEVER,
+)

mypyc/transform/log_trace.py

@@ -0,0 +1,84 @@
+"""This optional pass adds logging of various executed operations.
+
+Some subset of the executed operations are logged to the mypyc_trace.txt file.
+
+This is useful for performance analysis. For example, it's possible
+to identify how frequently various primitive functions are called,
+and in which code locations they are called.
+"""
+
+from __future__ import annotations
+
+from mypyc.ir.func_ir import FuncIR
+from mypyc.ir.ops import Call, CallC, CString, LoadLiteral, LoadStatic, Op, PrimitiveOp, Value
+from mypyc.irbuild.ll_builder import LowLevelIRBuilder
+from mypyc.options import CompilerOptions
+from mypyc.primitives.misc_ops import log_trace_event
+from mypyc.transform.ir_transform import IRTransform
+
+
+def insert_event_trace_logging(fn: FuncIR, options: CompilerOptions) -> None:
+    builder = LowLevelIRBuilder(None, options)
+    transform = LogTraceEventTransform(builder, fn.decl.fullname)
+    transform.transform_blocks(fn.blocks)
+    fn.blocks = builder.blocks
+
+
+def get_load_global_name(op: CallC) -> str | None:
+    name = op.function_name
+    if name == "CPyDict_GetItem":
+        arg = op.args[0]
+        if (
+            isinstance(arg, LoadStatic)
+            and arg.namespace == "static"
+            and arg.identifier == "globals"
+            and isinstance(op.args[1], LoadLiteral)
+        ):
+            return str(op.args[1].value)
+    return None
+
+
+class LogTraceEventTransform(IRTransform):
+    def __init__(self, builder: LowLevelIRBuilder, fullname: str) -> None:
+        super().__init__(builder)
+        self.fullname = fullname.encode("utf-8")
+
+    def visit_call(self, op: Call) -> Value:
+        # TODO: Use different op name when constructing an instance
+        return self.log(op, "call", op.fn.fullname)
+
+    def visit_primitive_op(self, op: PrimitiveOp) -> Value:
+        self.log(op, "primitive_op", op.desc.name)
+        return self.add(op)
+
+    def visit_call_c(self, op: CallC) -> Value:
+        if global_name := get_load_global_name(op):
+            return self.log(op, "globals_dict_get_item", global_name)
+
+        func_name = op.function_name
+        if func_name == "PyObject_Vectorcall" and isinstance(op.args[0], CallC):
+            if global_name := get_load_global_name(op.args[0]):
+                return self.log(op, "python_call_global", global_name)
+        elif func_name == "CPyObject_GetAttr" and isinstance(op.args[1], LoadLiteral):
+            return self.log(op, "python_get_attr", str(op.args[1].value))
+        elif func_name == "PyObject_VectorcallMethod" and isinstance(op.args[0], LoadLiteral):
+            return self.log(op, "python_call_method", str(op.args[0].value))
+
+        return self.log(op, "call_c", func_name)
+
+    def log(self, op: Op, name: str, details: str) -> Value:
+        if op.line >= 0:
+            line_str = str(op.line)
+        else:
+            line_str = ""
+        self.builder.primitive_op(
+            log_trace_event,
+            [
+                CString(self.fullname),
+                CString(line_str.encode("ascii")),
+                CString(name.encode("utf-8")),
+                CString(details.encode("utf-8")),
+            ],
+            op.line,
+        )
+        return self.add(op)