diff --git a/python/quadrants/lang/_func_base.py b/python/quadrants/lang/_func_base.py index aa94c98345..ba199fc36e 100644 --- a/python/quadrants/lang/_func_base.py +++ b/python/quadrants/lang/_func_base.py @@ -228,7 +228,7 @@ def __init__( self.arg_metas_expanded: list[ArgMetadata] = [] self.orig_arguments: list[ArgMetadata] = [] self.return_type = None - # Shared by every AST transform of this function; see ASTTransformerFuncContext.get_pos_info. + # Shared by every AST transform of this function; see ASTTransformerFuncContext.memoized_get_pos_info. self.pos_info_cache: dict[tuple, str] = {} self.check_parameter_annotations() diff --git a/python/quadrants/lang/ast/ast_transformer.py b/python/quadrants/lang/ast/ast_transformer.py index d152595093..875d1938ed 100644 --- a/python/quadrants/lang/ast/ast_transformer.py +++ b/python/quadrants/lang/ast/ast_transformer.py @@ -108,7 +108,7 @@ def build_Name(ctx: ASTTransformerFuncContext, node: ast.Name): # the flattened-name path bypasses ``build_Attribute`` entirely, so we must promote here too. node.ptr = ASTTransformer._promote_ndarray_if_declared(ctx, node.ptr) if isinstance(node, (ast.stmt, ast.expr)) and isinstance(node.ptr, Expr): - node.ptr.dbg_info = _qd_core.DebugInfo(ctx.get_pos_info(node)) + node.ptr.dbg_info = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) node.ptr.ptr.set_dbg_info(node.ptr.dbg_info) # ``qd.static`` is intentionally NOT a purity escape hatch: a captured module global is still flagged inside # a static scope, since its value never enters the fastcache key regardless of static wrapping. @@ -634,7 +634,7 @@ def build_Return(ctx: ASTTransformerFuncContext, node: ast.Return) -> None: raise QuadrantsSyntaxError("The return type is not supported now!") ctx.ast_builder.create_kernel_exprgroup_return( expr.make_expr_group(return_exprs), - _qd_core.DebugInfo(ctx.get_pos_info(node)), + _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)), ) else: ctx.return_data = node.value.ptr @@ -1143,7 +1143,7 @@ def build_range_for(ctx: ASTTransformerFuncContext, node: ast.For) -> None: begin = qd_ops.cast(expr.Expr(0), primitive_types.i32) end = qd_ops.cast(end_expr, primitive_types.i32) - for_di = _qd_core.DebugInfo(ctx.get_pos_info(node)) + for_di = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) ctx.ast_builder.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr, for_di) ctx.loop_depth += 1 build_stmts(ctx, node.body) @@ -1161,7 +1161,7 @@ def build_ndrange_for(ctx: ASTTransformerFuncContext, node: ast.For) -> None: primitive_types.i32, ) ndrange_loop_var = expr.Expr(ctx.ast_builder.make_id_expr("")) - for_di = _qd_core.DebugInfo(ctx.get_pos_info(node)) + for_di = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) ctx.ast_builder.begin_frontend_range_for(ndrange_loop_var.ptr, ndrange_begin.ptr, ndrange_end.ptr, for_di) I = impl.expr_init(ndrange_loop_var) targets = ASTTransformer.get_for_loop_targets(node) @@ -1214,7 +1214,7 @@ def build_grouped_ndrange_for(ctx: ASTTransformerFuncContext, node: ast.For) -> primitive_types.i32, ) ndrange_loop_var = expr.Expr(ctx.ast_builder.make_id_expr("")) - for_di = _qd_core.DebugInfo(ctx.get_pos_info(node)) + for_di = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) ctx.ast_builder.begin_frontend_range_for(ndrange_loop_var.ptr, ndrange_begin.ptr, ndrange_end.ptr, for_di) targets = ASTTransformer.get_for_loop_targets(node) @@ -1351,7 +1351,7 @@ def build_nested_mesh_for(ctx: ASTTransformerFuncContext, node: ast.For) -> None ctx.create_variable(loop_name, loop_var) begin = expr.Expr(0) end = qd_ops.cast(node.iter.ptr.size, primitive_types.i32) - for_di = _qd_core.DebugInfo(ctx.get_pos_info(node)) + for_di = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) ctx.ast_builder.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr, for_di) entry_expr = _qd_core.get_relation_access( ctx.mesh.mesh_ptr, @@ -1537,7 +1537,7 @@ def build_While(ctx: ASTTransformerFuncContext, node: ast.While) -> None: return None with ctx.loop_scope_guard(): - stmt_dbg_info = _qd_core.DebugInfo(ctx.get_pos_info(node)) + stmt_dbg_info = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) ctx.ast_builder.begin_frontend_while(expr.Expr(1, dtype=primitive_types.i32).ptr, stmt_dbg_info) while_cond = build_stmt(ctx, node.test) impl.begin_frontend_if(ctx.ast_builder, while_cond, stmt_dbg_info) @@ -1563,7 +1563,7 @@ def build_If(ctx: ASTTransformerFuncContext, node: ast.If) -> ast.If | None: return node with ctx.non_static_if_guard(node): - stmt_dbg_info = _qd_core.DebugInfo(ctx.get_pos_info(node)) + stmt_dbg_info = _qd_core.DebugInfo(ctx.memoized_get_pos_info(node)) impl.begin_frontend_if(ctx.ast_builder, node.test.ptr, stmt_dbg_info) ctx.ast_builder.begin_frontend_if_true() build_stmts(ctx, node.body) @@ -1684,7 +1684,7 @@ def build_Assert(ctx: ASTTransformerFuncContext, node: ast.Assert) -> None: else: msg = unparse(node.test) test = build_stmt(ctx, node.test) - impl.qd_assert(test, msg.strip(), extra_args, _qd_core.DebugInfo(ctx.get_pos_info(node))) + impl.qd_assert(test, msg.strip(), extra_args, _qd_core.DebugInfo(ctx.memoized_get_pos_info(node))) return None @staticmethod @@ -1692,7 +1692,7 @@ def build_Break(ctx: ASTTransformerFuncContext, node: ast.Break) -> None: if ctx.is_in_static_for(): nearest_non_static_if = ctx.current_loop_scope().nearest_non_static_if if nearest_non_static_if: - msg = ctx.get_pos_info(nearest_non_static_if.test) + msg = ctx.memoized_get_pos_info(nearest_non_static_if.test) msg += ( "You are trying to `break` a static `for` loop, " "but the `break` statement is inside a non-static `if`. " @@ -1700,7 +1700,7 @@ def build_Break(ctx: ASTTransformerFuncContext, node: ast.Break) -> None: raise QuadrantsSyntaxError(msg) ctx.set_loop_status(LoopStatus.Break) else: - ctx.ast_builder.insert_break_stmt(_qd_core.DebugInfo(ctx.get_pos_info(node))) + ctx.ast_builder.insert_break_stmt(_qd_core.DebugInfo(ctx.memoized_get_pos_info(node))) return None @staticmethod @@ -1708,7 +1708,7 @@ def build_Continue(ctx: ASTTransformerFuncContext, node: ast.Continue) -> None: if ctx.is_in_static_for(): nearest_non_static_if = ctx.current_loop_scope().nearest_non_static_if if nearest_non_static_if: - msg = ctx.get_pos_info(nearest_non_static_if.test) + msg = ctx.memoized_get_pos_info(nearest_non_static_if.test) msg += ( "You are trying to `continue` a static `for` loop, " "but the `continue` statement is inside a non-static `if`. " @@ -1716,7 +1716,7 @@ def build_Continue(ctx: ASTTransformerFuncContext, node: ast.Continue) -> None: raise QuadrantsSyntaxError(msg) ctx.set_loop_status(LoopStatus.Continue) else: - ctx.ast_builder.insert_continue_stmt(_qd_core.DebugInfo(ctx.get_pos_info(node))) + ctx.ast_builder.insert_continue_stmt(_qd_core.DebugInfo(ctx.memoized_get_pos_info(node))) return None @staticmethod diff --git a/python/quadrants/lang/ast/ast_transformer_utils.py b/python/quadrants/lang/ast/ast_transformer_utils.py index 0279393db7..34a8edce42 100644 --- a/python/quadrants/lang/ast/ast_transformer_utils.py +++ b/python/quadrants/lang/ast/ast_transformer_utils.py @@ -35,7 +35,7 @@ def __call__(self, ctx: "ASTTransformerFuncContext", node: ast.AST): if method is None: error_msg = f'Unsupported node "{node.__class__.__name__}"' raise QuadrantsSyntaxError(error_msg) - info = ctx.get_pos_info(node) if isinstance(node, (ast.stmt, ast.expr)) else "" + info = ctx.memoized_get_pos_info(node) if isinstance(node, (ast.stmt, ast.expr)) else "" with impl.get_runtime().src_info_guard(info): res = method(ctx, node) if not hasattr(node, "violates_pure"): @@ -52,7 +52,7 @@ def __call__(self, ctx: "ASTTransformerFuncContext", node: ast.AST): ctx.raised = True e = handle_exception_from_cpp(e) if not isinstance(e, QuadrantsCompilationError): - msg = ctx.get_pos_info(node) + traceback.format_exc() + msg = ctx.memoized_get_pos_info(node) + traceback.format_exc() raise QuadrantsCompilationError(msg) from None msg = f"""quadrants stack trace: === @@ -60,7 +60,7 @@ def __call__(self, ctx: "ASTTransformerFuncContext", node: ast.AST): === Your code: -{ctx.get_pos_info(node)}{e} +{ctx.memoized_get_pos_info(node)}{e} """ raise type(e)(msg) from None @@ -414,17 +414,9 @@ def get_var_by_name(self, name: str) -> tuple[bool, Any, str | None]: except AttributeError: raise QuadrantsNameError(f'Name "{name}" is not defined') - def get_pos_info(self, node: ast.AST) -> str: - # Runs for every stmt/expr node of every transform (see ASTTransformerBase.__call__), and the TextWrapper - # formatting below dominates the Python side of a kernel build. The same function is transformed once per - # inlined call site -- tens of times over -- and each transform re-formats the identical positions, so - # memoise on the function rather than on this context, which exists for a single transform. - # - # The cache lives on the FuncBase because that is what fixes everything the result depends on beyond the - # node: `file`, `src`, `indent`, `lineno_offset` and the function name all derive from it and are identical - # across its transforms. Scoping it there also keeps it bounded by the function's lifetime and correct - # across a module reload, which hands out new FuncBase objects and so a new cache. + def memoized_get_pos_info(self, node: ast.AST) -> str: key = (node.lineno, node.col_offset, node.end_lineno, node.end_col_offset, node.__class__.__name__) + # self.func is persistent, whereas context is ephemeral cached = self.func.pos_info_cache.get(key) if cached is not None: return cached @@ -433,6 +425,11 @@ def get_pos_info(self, node: ast.AST) -> str: return msg def _build_pos_info(self, node: ast.AST) -> str: + """ + Returns a formatted string: a source-location snippet for node, consisting of a location header, the relevant + source line(s), and a caret (^) underline marking the exact columns the node spans — the same shape as a + Python traceback frame. + """ msg = f'File "{self.file}", line {node.lineno + self.lineno_offset}, in {self.func.func.__name__}:\n' col_offset = self.indent + node.col_offset end_col_offset = self.indent + node.end_col_offset