From f51e0b1be40b970dd934c1f8c98c04c9bb3eba63 Mon Sep 17 00:00:00 2001 From: Keren Zhou Date: Wed, 23 Mar 2022 00:09:49 -0500 Subject: [PATCH] [FRONTEND] Hot fix for lineno (#481) Override __reduce__ to make CompilationError pickable and print out error messages --- python/triton/code_gen.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/python/triton/code_gen.py b/python/triton/code_gen.py index bb781dfe6..13821f8ca 100644 --- a/python/triton/code_gen.py +++ b/python/triton/code_gen.py @@ -677,9 +677,13 @@ class CompilationError(Exception): self.message = f'at {node.lineno}:{node.col_offset}:\n' self.message += '\n'.join(src.split('\n')[:node.lineno]) self.message += '\n' + ' ' * node.col_offset + '^' + self.src = src + self.node = node super().__init__(self.message) + + def __reduce__(self): # this is necessary to make CompilationError picklable - self.args = (src, node) + return (type(self), (self.src, self.node)) class OutOfResources(Exception): @@ -687,8 +691,14 @@ class OutOfResources(Exception): self.message = f'out of resource: {name}, '\ f'Required: {required}, '\ f'Hardware limit: {limit}' + self.required = required + self.limit = limit + self.name = name super().__init__(self.message) - self.args = (required, limit, name) + + def __reduce__(self): + # this is necessary to make CompilationError picklable + return (type(self), (self.required, self.limit, self.name)) class Kernel: