Store the generated source code in jit._Handle for debugging purposes

InfiniTensor · Oct 21, 2024 · b3edc8d · b3edc8d
1 parent 34d3e04
commit b3edc8d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/ninetoothed/jit.py b/src/ninetoothed/jit.py
@@ -59,6 +59,7 @@ def __call__(self):
         handle = _Handle(
             namespace[self.func.__name__],
             namespace[f"launch_{self.func.__name__}"],
+            unparsed,
         )
 
         type(self).handles[source_file][source_line] = handle
@@ -459,9 +460,10 @@ def visit_Call(self, node):
 
 
 class _Handle:
-    def __init__(self, kernel, launch):
+    def __init__(self, kernel, launch, source):
         self._kernel = kernel
         self._launch = launch
+        self._source = source
 
     def __call__(self, *args, **kwargs):
         return self._launch(*args, **kwargs)