warp-lang 1.6.1__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/codegen.py

@@ -1,9 +1,17 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
 
 from __future__ import annotations
 
@@ -18,7 +26,7 @@ import re
 import sys
 import textwrap
 import types
-from typing import Any, Callable, Dict, Mapping, Optional, Sequence
+from typing import Any, Callable, Dict, Mapping, Optional, Sequence, get_args, get_origin
 
 import warp.config
 from warp.types import *
@@ -49,7 +57,7 @@ class WarpCodegenKeyError(KeyError):
 
 
 # map operator to function name
-builtin_operators = {}
+builtin_operators: Dict[type[ast.AST], str] = {}
 
 # see https://www.ics.uci.edu/~pattis/ICS-31/lectures/opexp.pdf for a
 # nice overview of python operators
@@ -114,16 +122,6 @@ def get_closure_cell_contents(obj):
     return None
 
 
-def get_type_origin(tp):
-    # Compatible version of `typing.get_origin()` for Python 3.7 and older.
-    return getattr(tp, "__origin__", None)
-
-
-def get_type_args(tp):
-    # Compatible version of `typing.get_args()` for Python 3.7 and older.
-    return getattr(tp, "__args__", ())
-
-
 def eval_annotations(annotations: Mapping[str, Any], obj: Any) -> Mapping[str, Any]:
     """Un-stringize annotations caused by `from __future__ import annotations` of PEP 563."""
     # Implementation backported from `inspect.get_annotations()` for Python 3.9 and older.
@@ -407,12 +405,14 @@ class StructInstance:
 
 
 class Struct:
-    def __init__(self, cls, key, module):
+    hash: bytes
+
+    def __init__(self, cls: type, key: str, module: warp.context.Module):
         self.cls = cls
         self.module = module
         self.key = key
+        self.vars: Dict[str, Var] = {}
 
-        self.vars = {}
         annotations = get_annotations(self.cls)
         for label, type in annotations.items():
             self.vars[label] = Var(label, type)
@@ -583,11 +583,11 @@ class Reference:
         self.value_type = value_type
 
 
-def is_reference(type):
+def is_reference(type: Any) -> builtins.bool:
     return isinstance(type, Reference)
 
 
-def strip_reference(arg):
+def strip_reference(arg: Any) -> Any:
     if is_reference(arg):
         return arg.value_type
     else:
@@ -615,7 +615,15 @@ def compute_type_str(base_name, template_params):
 
 
 class Var:
-    def __init__(self, label, type, requires_grad=False, constant=None, prefix=True):
+    def __init__(
+        self,
+        label: str,
+        type: type,
+        requires_grad: builtins.bool = False,
+        constant: Optional[builtins.bool] = None,
+        prefix: builtins.bool = True,
+        relative_lineno: Optional[int] = None,
+    ):
         # convert built-in types to wp types
         if type == float:
             type = float32
@@ -638,11 +646,14 @@ class Var:
         # used to associate a view array Var with its parent array Var
         self.parent = None
 
+        # Used to associate the variable with the Python statement that resulted in it being created.
+        self.relative_lineno = relative_lineno
+
     def __str__(self):
         return self.label
 
     @staticmethod
-    def type_to_ctype(t, value_type=False):
+    def type_to_ctype(t: type, value_type: builtins.bool = False) -> str:
         if is_array(t):
             if hasattr(t.dtype, "_wp_generic_type_str_"):
                 dtypestr = compute_type_str(f"wp::{t.dtype._wp_generic_type_str_}", t.dtype._wp_type_params_)
@@ -673,7 +684,7 @@ class Var:
         else:
             return f"wp::{t.__name__}"
 
-    def ctype(self, value_type=False):
+    def ctype(self, value_type: builtins.bool = False) -> str:
         return Var.type_to_ctype(self.type, value_type)
 
     def emit(self, prefix: str = "var"):
@@ -795,7 +806,7 @@ def func_match_args(func, arg_types, kwarg_types):
     return True
 
 
-def get_arg_type(arg: Union[Var, Any]):
+def get_arg_type(arg: Union[Var, Any]) -> type:
     if isinstance(arg, str):
         return str
 
@@ -811,7 +822,7 @@ def get_arg_type(arg: Union[Var, Any]):
     return type(arg)
 
 
-def get_arg_value(arg: Union[Var, Any]):
+def get_arg_value(arg: Any) -> Any:
     if isinstance(arg, Sequence):
         return tuple(get_arg_value(x) for x in arg)
 
@@ -859,6 +870,9 @@ class Adjoint:
                 "please save it on a file and use `importlib` if needed."
             ) from e
 
+        # Indicates where the function definition starts (excludes decorators)
+        adj.fun_def_lineno = None
+
         # get function source code
         adj.source = inspect.getsource(func)
         # ensures that indented class methods can be parsed as kernels
@@ -933,9 +947,6 @@ class Adjoint:
         # for unit testing errors being spit out from kernels.
         adj.skip_build = False
 
-        # Collect the LTOIR required at link-time
-        adj.ltoirs = []
-
     # allocate extra space for a function call that requires its
     # own shared memory space, we treat shared memory as a stack
     # where each function pushes and pops space off, the extra
@@ -1125,7 +1136,7 @@ class Adjoint:
         name = str(index)
 
         # allocate new variable
-        v = Var(name, type=type, constant=constant)
+        v = Var(name, type=type, constant=constant, relative_lineno=adj.lineno)
 
         adj.variables.append(v)
 
@@ -1150,11 +1161,44 @@ class Adjoint:
 
         return var
 
-    # append a statement to the forward pass
-    def add_forward(adj, statement, replay=None, skip_replay=False):
+    def get_line_directive(adj, statement: str, relative_lineno: Optional[int] = None) -> Optional[str]:
+        """Get a line directive for the given statement.
+
+        Args:
+            statement: The statement to get the line directive for.
+            relative_lineno: The line number of the statement relative to the function.
+
+        Returns:
+            A line directive for the given statement, or None if no line directive is needed.
+        """
+
+        # lineinfo is enabled by default in debug mode regardless of the builder option, don't want to unnecessarily
+        # emit line directives in generated code if it's not being compiled with line information
+        lineinfo_enabled = (
+            adj.builder_options.get("lineinfo", False) or adj.builder_options.get("mode", "release") == "debug"
+        )
+
+        if relative_lineno is not None and lineinfo_enabled and warp.config.line_directives:
+            is_comment = statement.strip().startswith("//")
+            if not is_comment:
+                line = relative_lineno + adj.fun_lineno
+                # Convert backslashes to forward slashes for CUDA compatibility
+                normalized_path = adj.filename.replace("\\", "/")
+                return f'#line {line} "{normalized_path}"'
+        return None
+
+    def add_forward(adj, statement: str, replay: Optional[str] = None, skip_replay: builtins.bool = False) -> None:
+        """Append a statement to the forward pass."""
+
+        if line_directive := adj.get_line_directive(statement, adj.lineno):
+            adj.blocks[-1].body_forward.append(line_directive)
+
         adj.blocks[-1].body_forward.append(adj.indentation + statement)
 
         if not skip_replay:
+            if line_directive:
+                adj.blocks[-1].body_replay.append(line_directive)
+
             if replay:
                 # if custom replay specified then output it
                 adj.blocks[-1].body_replay.append(adj.indentation + replay)
@@ -1163,9 +1207,14 @@ class Adjoint:
                 adj.blocks[-1].body_replay.append(adj.indentation + statement)
 
     # append a statement to the reverse pass
-    def add_reverse(adj, statement):
+    def add_reverse(adj, statement: str) -> None:
+        """Append a statement to the reverse pass."""
+
         adj.blocks[-1].body_reverse.append(adj.indentation + statement)
 
+        if line_directive := adj.get_line_directive(statement, adj.lineno):
+            adj.blocks[-1].body_reverse.append(line_directive)
+
     def add_constant(adj, n):
         output = adj.add_var(type=type(n), constant=n)
         return output
@@ -1273,7 +1322,7 @@ class Adjoint:
 
         # Bind the positional and keyword arguments to the function's signature
         # in order to process them as Python does it.
-        bound_args = func.signature.bind(*args, **kwargs)
+        bound_args: inspect.BoundArguments = func.signature.bind(*args, **kwargs)
 
         # Type args are the “compile time” argument values we get from codegen.
         # For example, when calling `wp.vec3f(...)` from within a kernel,
@@ -1616,6 +1665,8 @@ class Adjoint:
         adj.blocks[-1].body_reverse.extend(reversed(reverse))
 
     def emit_FunctionDef(adj, node):
+        adj.fun_def_lineno = node.lineno
+
         for f in node.body:
             # Skip variable creation for standalone constants, including docstrings
             if isinstance(f, ast.Expr) and isinstance(f.value, ast.Constant):
@@ -1680,7 +1731,7 @@ class Adjoint:
 
             if var1 != var2:
                 # insert a phi function that selects var1, var2 based on cond
-                out = adj.add_builtin_call("select", [cond, var1, var2])
+                out = adj.add_builtin_call("where", [cond, var2, var1])
                 adj.symbols[sym] = out
 
         symbols_prev = adj.symbols.copy()
@@ -1704,7 +1755,7 @@ class Adjoint:
             if var1 != var2:
                 # insert a phi function that selects var1, var2 based on cond
                 # note the reversed order of vars since we want to use !cond as our select
-                out = adj.add_builtin_call("select", [cond, var2, var1])
+                out = adj.add_builtin_call("where", [cond, var1, var2])
                 adj.symbols[sym] = out
 
     def emit_Compare(adj, node):
@@ -1848,25 +1899,6 @@ class Adjoint:
                 ) from e
             raise WarpCodegenAttributeError(f"Error, `{node.attr}` is not an attribute of '{aggregate}'") from e
 
-    def emit_String(adj, node):
-        # string constant
-        return adj.add_constant(node.s)
-
-    def emit_Num(adj, node):
-        # lookup constant, if it has already been assigned then return existing var
-        key = (node.n, type(node.n))
-
-        if key in adj.symbols:
-            return adj.symbols[key]
-        else:
-            out = adj.add_constant(node.n)
-            adj.symbols[key] = out
-            return out
-
-    def emit_Ellipsis(adj, node):
-        # stubbed @wp.native_func
-        return
-
     def emit_Assert(adj, node):
         # eval condition
         cond = adj.eval(node.test)
@@ -1878,24 +1910,11 @@ class Adjoint:
 
         adj.add_forward(f'assert(("{escaped_segment}",{cond.emit()}));')
 
-    def emit_NameConstant(adj, node):
-        if node.value:
-            return adj.add_constant(node.value)
-        elif node.value is None:
-            raise WarpCodegenTypeError("None type unsupported")
-        else:
-            return adj.add_constant(False)
-
     def emit_Constant(adj, node):
-        if isinstance(node, ast.Str):
-            return adj.emit_String(node)
-        elif isinstance(node, ast.Num):
-            return adj.emit_Num(node)
-        elif isinstance(node, ast.Ellipsis):
-            return adj.emit_Ellipsis(node)
+        if node.value is None:
+            raise WarpCodegenTypeError("None type unsupported")
         else:
-            assert isinstance(node, ast.NameConstant) or isinstance(node, ast.Constant)
-            return adj.emit_NameConstant(node)
+            return adj.add_constant(node.value)
 
     def emit_BinOp(adj, node):
         # evaluate binary operator arguments
@@ -1989,10 +2008,11 @@ class Adjoint:
         adj.end_while()
 
     def eval_num(adj, a):
-        if isinstance(a, ast.Num):
-            return True, a.n
-        if isinstance(a, ast.UnaryOp) and isinstance(a.op, ast.USub) and isinstance(a.operand, ast.Num):
-            return True, -a.operand.n
+        if isinstance(a, ast.Constant):
+            return True, a.value
+        if isinstance(a, ast.UnaryOp) and isinstance(a.op, ast.USub) and isinstance(a.operand, ast.Constant):
+            # Negative constant
+            return True, -a.operand.value
 
         # try and resolve the expression to an object
         # e.g.: wp.constant in the globals scope
@@ -2522,8 +2542,8 @@ class Adjoint:
                             f"Warning: mutating {node_source} in function {adj.fun_name} at {adj.filename}:{lineno}: this is a non-differentiable operation.\n{line}\n"
                         )
                 else:
-                    if adj.builder_options.get("enable_backward", True):
-                        out = adj.add_builtin_call("assign", [target, *indices, rhs])
+                    if warp.config.enable_vector_component_overwrites:
+                        out = adj.add_builtin_call("assign_copy", [target, *indices, rhs])
 
                         # re-point target symbol to out var
                         for id in adj.symbols:
@@ -2531,8 +2551,7 @@ class Adjoint:
                                 adj.symbols[id] = out
                                 break
                     else:
-                        attr = adj.add_builtin_call("index", [target, *indices])
-                        adj.add_builtin_call("store", [attr, rhs])
+                        adj.add_builtin_call("assign_inplace", [target, *indices, rhs])
 
             else:
                 raise WarpCodegenError(
@@ -2575,8 +2594,8 @@ class Adjoint:
                     attr = adj.add_builtin_call("indexref", [aggregate, index])
                     adj.add_builtin_call("store", [attr, rhs])
                 else:
-                    if adj.builder_options.get("enable_backward", True):
-                        out = adj.add_builtin_call("assign", [aggregate, index, rhs])
+                    if warp.config.enable_vector_component_overwrites:
+                        out = adj.add_builtin_call("assign_copy", [aggregate, index, rhs])
 
                         # re-point target symbol to out var
                         for id in adj.symbols:
@@ -2584,8 +2603,7 @@ class Adjoint:
                                 adj.symbols[id] = out
                                 break
                     else:
-                        attr = adj.add_builtin_call("index", [aggregate, index])
-                        adj.add_builtin_call("store", [attr, rhs])
+                        adj.add_builtin_call("assign_inplace", [aggregate, index, rhs])
 
             else:
                 attr = adj.emit_Attribute(lhs)
@@ -2691,10 +2709,12 @@ class Adjoint:
 
             elif type_is_vector(target_type) or type_is_quaternion(target_type) or type_is_matrix(target_type):
                 if isinstance(node.op, ast.Add):
-                    adj.add_builtin_call("augassign_add", [target, *indices, rhs])
+                    adj.add_builtin_call("add_inplace", [target, *indices, rhs])
                 elif isinstance(node.op, ast.Sub):
-                    adj.add_builtin_call("augassign_sub", [target, *indices, rhs])
+                    adj.add_builtin_call("sub_inplace", [target, *indices, rhs])
                 else:
+                    if warp.config.verbose:
+                        print(f"Warning: in-place op {node.op} is not differentiable")
                     make_new_assign_statement()
                     return
 
@@ -2724,9 +2744,6 @@ class Adjoint:
         ast.BoolOp: emit_BoolOp,
         ast.Name: emit_Name,
         ast.Attribute: emit_Attribute,
-        ast.Str: emit_String,  # Deprecated in 3.8; use Constant
-        ast.Num: emit_Num,  # Deprecated in 3.8; use Constant
-        ast.NameConstant: emit_NameConstant,  # Deprecated in 3.8; use Constant
         ast.Constant: emit_Constant,
         ast.BinOp: emit_BinOp,
         ast.UnaryOp: emit_UnaryOp,
@@ -2736,14 +2753,13 @@ class Adjoint:
         ast.Continue: emit_Continue,
         ast.Expr: emit_Expr,
         ast.Call: emit_Call,
-        ast.Index: emit_Index,  # Deprecated in 3.8; Use the index value directly instead.
+        ast.Index: emit_Index,  # Deprecated in 3.9
         ast.Subscript: emit_Subscript,
         ast.Assign: emit_Assign,
         ast.Return: emit_Return,
         ast.AugAssign: emit_AugAssign,
         ast.Tuple: emit_Tuple,
         ast.Pass: emit_Pass,
-        ast.Ellipsis: emit_Ellipsis,
         ast.Assert: emit_Assert,
     }
 
@@ -2939,12 +2955,16 @@ class Adjoint:
 
             # We want to replace the expression code in-place,
             # so reparse it to get the correct column info.
-            len_value_locs = []
+            len_value_locs: List[Tuple[int, int, int]] = []
             expr_tree = ast.parse(static_code)
             assert len(expr_tree.body) == 1 and isinstance(expr_tree.body[0], ast.Expr)
             expr_root = expr_tree.body[0].value
             for expr_node in ast.walk(expr_root):
-                if isinstance(expr_node, ast.Call) and expr_node.func.id == "len" and len(expr_node.args) == 1:
+                if (
+                    isinstance(expr_node, ast.Call)
+                    and getattr(expr_node.func, "id", None) == "len"
+                    and len(expr_node.args) == 1
+                ):
                     len_expr = static_code[expr_node.col_offset : expr_node.end_col_offset]
                     try:
                         len_value = eval(len_expr, len_expr_ctx)
@@ -3102,9 +3122,9 @@ class Adjoint:
 
         local_variables = set()  # Track local variables appearing on the LHS so we know when variables are shadowed
 
-        constants = {}
-        types = {}
-        functions = {}
+        constants: Dict[str, Any] = {}
+        types: Dict[Union[Struct, type], Any] = {}
+        functions: Dict[warp.context.Function, Any] = {}
 
         for node in ast.walk(adj.tree):
             if isinstance(node, ast.Name) and node.id not in local_variables:
@@ -3147,7 +3167,7 @@ class Adjoint:
 # code generation
 
 cpu_module_header = """
-#define WP_TILE_BLOCK_DIM {tile_size}
+#define WP_TILE_BLOCK_DIM {block_dim}
 #define WP_NO_CRT
 #include "builtin.h"
 
@@ -3166,7 +3186,7 @@ cpu_module_header = """
 """
 
 cuda_module_header = """
-#define WP_TILE_BLOCK_DIM {tile_size}
+#define WP_TILE_BLOCK_DIM {block_dim}
 #define WP_NO_CRT
 #include "builtin.h"
 
@@ -3189,6 +3209,7 @@ struct {name}
 {{
 {struct_body}
 
+    {defaulted_constructor_def}
     CUDA_CALLABLE {name}({forward_args})
     {forward_initializers}
     {{
@@ -3231,53 +3252,53 @@ static void adj_{name}(
 
 cuda_forward_function_template = """
 // {filename}:{lineno}
-static CUDA_CALLABLE {return_type} {name}(
+{line_directive}static CUDA_CALLABLE {return_type} {name}(
     {forward_args})
 {{
-{forward_body}}}
+{forward_body}{line_directive}}}
 
 """
 
 cuda_reverse_function_template = """
 // {filename}:{lineno}
-static CUDA_CALLABLE void adj_{name}(
+{line_directive}static CUDA_CALLABLE void adj_{name}(
     {reverse_args})
 {{
-{reverse_body}}}
+{reverse_body}{line_directive}}}
 
 """
 
 cuda_kernel_template_forward = """
 
-extern "C" __global__ void {name}_cuda_kernel_forward(
+{line_directive}extern "C" __global__ void {name}_cuda_kernel_forward(
     {forward_args})
 {{
-    for (size_t _idx = static_cast<size_t>(blockDim.x) * static_cast<size_t>(blockIdx.x) + static_cast<size_t>(threadIdx.x);
-         _idx < dim.size;
-         _idx += static_cast<size_t>(blockDim.x) * static_cast<size_t>(gridDim.x))
+{line_directive}    for (size_t _idx = static_cast<size_t>(blockDim.x) * static_cast<size_t>(blockIdx.x) + static_cast<size_t>(threadIdx.x);
+{line_directive}         _idx < dim.size;
+{line_directive}         _idx += static_cast<size_t>(blockDim.x) * static_cast<size_t>(gridDim.x))
     {{
         // reset shared memory allocator
-        wp::tile_alloc_shared(0, true);
+{line_directive}        wp::tile_alloc_shared(0, true);
 
-{forward_body}    }}
-}}
+{forward_body}{line_directive}    }}
+{line_directive}}}
 
 """
 
 cuda_kernel_template_backward = """
 
-extern "C" __global__ void {name}_cuda_kernel_backward(
+{line_directive}extern "C" __global__ void {name}_cuda_kernel_backward(
     {reverse_args})
 {{
-    for (size_t _idx = static_cast<size_t>(blockDim.x) * static_cast<size_t>(blockIdx.x) + static_cast<size_t>(threadIdx.x);
-         _idx < dim.size;
-         _idx += static_cast<size_t>(blockDim.x) * static_cast<size_t>(gridDim.x))
+{line_directive}    for (size_t _idx = static_cast<size_t>(blockDim.x) * static_cast<size_t>(blockIdx.x) + static_cast<size_t>(threadIdx.x);
+{line_directive}         _idx < dim.size;
+{line_directive}         _idx += static_cast<size_t>(blockDim.x) * static_cast<size_t>(gridDim.x))
     {{
         // reset shared memory allocator
-        wp::tile_alloc_shared(0, true);
+{line_directive}        wp::tile_alloc_shared(0, true);
 
-{reverse_body}    }}
-}}
+{reverse_body}{line_directive}    }}
+{line_directive}}}
 
 """
 
@@ -3307,10 +3328,17 @@ extern "C" {{
 WP_API void {name}_cpu_forward(
     {forward_args})
 {{
-    for (size_t task_index = 0; task_index < dim.size; ++task_index)
+for (size_t task_index = 0; task_index < dim.size; ++task_index)
     {{
+        // init shared memory allocator
+        wp::tile_alloc_shared(0, true);
+
         {name}_cpu_kernel_forward(
             {forward_params});
+
+        // check shared memory allocator
+        wp::tile_alloc_shared(0, false, true);
+
     }}
 }}
 
@@ -3327,8 +3355,14 @@ WP_API void {name}_cpu_backward(
 {{
     for (size_t task_index = 0; task_index < dim.size; ++task_index)
     {{
+        // initialize shared memory allocator
+        wp::tile_alloc_shared(0, true);
+
         {name}_cpu_kernel_backward(
             {reverse_params});
+
+        // check shared memory allocator
+        wp::tile_alloc_shared(0, false, true);
     }}
 }}
 
@@ -3410,7 +3444,7 @@ def indent(args, stops=1):
 
 
 # generates a C function name based on the python function name
-def make_full_qualified_name(func):
+def make_full_qualified_name(func: Union[str, Callable]) -> str:
     if not isinstance(func, str):
         func = func.__qualname__
     return re.sub("[^0-9a-zA-Z_]+", "", func.replace(".", "__"))
@@ -3440,7 +3474,8 @@ def codegen_struct(struct, device="cpu", indent_size=4):
     # forward args
     for label, var in struct.vars.items():
         var_ctype = var.ctype()
-        forward_args.append(f"{var_ctype} const& {label} = {{}}")
+        default_arg_def = " = {}" if forward_args else ""
+        forward_args.append(f"{var_ctype} const& {label}{default_arg_def}")
         reverse_args.append(f"{var_ctype} const&")
 
         namespace = "wp::" if var_ctype.startswith("wp::") or var_ctype == "bool" else ""
@@ -3464,6 +3499,9 @@ def codegen_struct(struct, device="cpu", indent_size=4):
 
     reverse_args.append(name + " & adj_ret")
 
+    # explicitly defaulted default constructor if no default constructor has been defined
+    defaulted_constructor_def = f"{name}() = default;" if forward_args else ""
+
     return struct_template.format(
         name=name,
         struct_body="".join([indent_block + l for l in body]),
@@ -3473,6 +3511,7 @@ def codegen_struct(struct, device="cpu", indent_size=4):
         reverse_body="".join(reverse_body),
         prefix_add_body="".join(prefix_add_body),
         atomic_add_body="".join(atomic_add_body),
+        defaulted_constructor_def=defaulted_constructor_def,
     )
 
 
@@ -3502,6 +3541,9 @@ def codegen_func_forward(adj, func_type="kernel", device="cpu"):
         else:
             lines += [f"const {var.ctype()} {var.emit()} = {constant_str(var.constant)};\n"]
 
+        if line_directive := adj.get_line_directive(lines[-1], var.relative_lineno):
+            lines.insert(-1, f"{line_directive}\n")
+
     # forward pass
     lines += ["//---------\n"]
     lines += ["// forward\n"]
@@ -3509,7 +3551,7 @@ def codegen_func_forward(adj, func_type="kernel", device="cpu"):
     for f in adj.blocks[0].body_forward:
         lines += [f + "\n"]
 
-    return "".join([indent_block + l for l in lines])
+    return "".join(l.lstrip() if l.lstrip().startswith("#line") else indent_block + l for l in lines)
 
 
 def codegen_func_reverse(adj, func_type="kernel", device="cpu"):
@@ -3539,6 +3581,9 @@ def codegen_func_reverse(adj, func_type="kernel", device="cpu"):
         else:
             lines += [f"const {var.ctype()} {var.emit()} = {constant_str(var.constant)};\n"]
 
+        if line_directive := adj.get_line_directive(lines[-1], var.relative_lineno):
+            lines.insert(-1, f"{line_directive}\n")
+
     # dual vars
     lines += ["//---------\n"]
     lines += ["// dual vars\n"]
@@ -3559,6 +3604,9 @@ def codegen_func_reverse(adj, func_type="kernel", device="cpu"):
         else:
             lines += [f"{ctype} {name} = {{}};\n"]
 
+        if line_directive := adj.get_line_directive(lines[-1], var.relative_lineno):
+            lines.insert(-1, f"{line_directive}\n")
+
     # forward pass
     lines += ["//---------\n"]
     lines += ["// forward\n"]
@@ -3579,7 +3627,7 @@ def codegen_func_reverse(adj, func_type="kernel", device="cpu"):
     else:
         lines += ["return;\n"]
 
-    return "".join([indent_block + l for l in lines])
+    return "".join(l.lstrip() if l.lstrip().startswith("#line") else indent_block + l for l in lines)
 
 
 def codegen_func(adj, c_func_name: str, device="cpu", options=None):
@@ -3587,11 +3635,11 @@ def codegen_func(adj, c_func_name: str, device="cpu", options=None):
         options = {}
 
     if adj.return_var is not None and "return" in adj.arg_types:
-        if get_type_origin(adj.arg_types["return"]) is tuple:
-            if len(get_type_args(adj.arg_types["return"])) != len(adj.return_var):
+        if get_origin(adj.arg_types["return"]) is tuple:
+            if len(get_args(adj.arg_types["return"])) != len(adj.return_var):
                 raise WarpCodegenError(
                     f"The function `{adj.fun_name}` has its return type "
-                    f"annotated as a tuple of {len(get_type_args(adj.arg_types['return']))} elements "
+                    f"annotated as a tuple of {len(get_args(adj.arg_types['return']))} elements "
                     f"but the code returns {len(adj.return_var)} values."
                 )
             elif not types_equal(adj.arg_types["return"], tuple(x.type for x in adj.return_var)):
@@ -3600,7 +3648,7 @@ def codegen_func(adj, c_func_name: str, device="cpu", options=None):
                     f"annotated as `{warp.context.type_str(adj.arg_types['return'])}` "
                     f"but the code returns a tuple with types `({', '.join(warp.context.type_str(x.type) for x in adj.return_var)})`."
                 )
-        elif len(adj.return_var) > 1 and get_type_origin(adj.arg_types["return"]) is not tuple:
+        elif len(adj.return_var) > 1 and get_origin(adj.arg_types["return"]) is not tuple:
             raise WarpCodegenError(
                 f"The function `{adj.fun_name}` has its return type "
                 f"annotated as `{warp.context.type_str(adj.arg_types['return'])}` "
@@ -3613,6 +3661,13 @@ def codegen_func(adj, c_func_name: str, device="cpu", options=None):
                 f"but the code returns a value of type `{warp.context.type_str(adj.return_var[0].type)}`."
             )
 
+    # Build line directive for function definition (subtract 1 to account for 1-indexing of AST line numbers)
+    # This is used as a catch-all C-to-Python source line mapping for any code that does not have
+    # a direct mapping to a Python source line.
+    func_line_directive = ""
+    if line_directive := adj.get_line_directive("", adj.fun_def_lineno - 1):
+        func_line_directive = f"{line_directive}\n"
+
     # forward header
     if adj.return_var is not None and len(adj.return_var) == 1:
         return_type = adj.return_var[0].ctype()
@@ -3676,6 +3731,7 @@ def codegen_func(adj, c_func_name: str, device="cpu", options=None):
             forward_body=forward_body,
             filename=adj.filename,
             lineno=adj.fun_lineno,
+            line_directive=func_line_directive,
         )
 
     if not adj.skip_reverse_codegen:
@@ -3694,6 +3750,7 @@ def codegen_func(adj, c_func_name: str, device="cpu", options=None):
             reverse_body=reverse_body,
             filename=adj.filename,
             lineno=adj.fun_lineno,
+            line_directive=func_line_directive,
         )
 
     return s
@@ -3736,6 +3793,7 @@ def codegen_snippet(adj, name, snippet, adj_snippet, replay_snippet):
         forward_body=snippet,
         filename=adj.filename,
         lineno=adj.fun_lineno,
+        line_directive="",
     )
 
     if replay_snippet is not None:
@@ -3746,6 +3804,7 @@ def codegen_snippet(adj, name, snippet, adj_snippet, replay_snippet):
             forward_body=replay_snippet,
             filename=adj.filename,
             lineno=adj.fun_lineno,
+            line_directive="",
         )
 
     if adj_snippet:
@@ -3761,6 +3820,7 @@ def codegen_snippet(adj, name, snippet, adj_snippet, replay_snippet):
         reverse_body=reverse_body,
         filename=adj.filename,
         lineno=adj.fun_lineno,
+        line_directive="",
     )
 
     return s
@@ -3773,6 +3833,13 @@ def codegen_kernel(kernel, device, options):
 
     adj = kernel.adj
 
+    # Build line directive for function definition (subtract 1 to account for 1-indexing of AST line numbers)
+    # This is used as a catch-all C-to-Python source line mapping for any code that does not have
+    # a direct mapping to a Python source line.
+    func_line_directive = ""
+    if line_directive := adj.get_line_directive("", adj.fun_def_lineno - 1):
+        func_line_directive = f"{line_directive}\n"
+
     if device == "cpu":
         template_forward = cpu_kernel_template_forward
         template_backward = cpu_kernel_template_backward
@@ -3800,6 +3867,7 @@ def codegen_kernel(kernel, device, options):
         {
             "forward_args": indent(forward_args),
             "forward_body": forward_body,
+            "line_directive": func_line_directive,
         }
     )
     template += template_forward