[frontend] Improve function def parsing with stacked decorators. #3564
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
zhxchen17
changed the title
jlebar
reviewed
Apr 4, 2024
python/triton/runtime/jit.py
Outdated
| @@ -478,7 +479,7 @@ def __init__(self, fn, version=None, do_not_specialize=None, debug=None, noinlin | |||
|
|
|||
| # function source code (without decorators) | |||
| self.src = textwrap.dedent(inspect.getsource(fn)) | |||
| self.src = self.src[self.src.find("def"):] | |||
| self.src = self.src[re.search(r"def (\w+)\s*\((.*?)\):", self.src).start():] | |||
There was a problem hiding this comment.
Can we remove the parens around (.*?)?
Also if we're not in multiline mode I think this does not work, because . only matches on the current line?
There was a problem hiding this comment.
yeah you're right, sorry I forgot the multiline case.
jlebar
reviewed
Apr 4, 2024
|
Addressed comments to minimize the test. Tweak the regex a little bit so that we always match against the beginning of a line, which is simpler to reason about. |
jlebar
approved these changes
Apr 4, 2024
Under certain conditions where we apply another decorator on top of
triton.jit, there is a small chance that in the decorator body there's
a "def" substring in it. This is problematic when JITFunction parses
function def using str.find('def'). This causes an issue compiling some
pytorch model when we happened to decorate a jit function with
a hash string argument which has "def" in it.
This diff tries a bit harder to locate the real "def" position using
a more complex regex which should further reduce the chance to hit this
error.
auto-merge was automatically disabled
April 4, 2024 02:33
Head branch was pushed to by a user without write access
|
Re-pushed this to fix yapf lint error. |
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
May 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
Jul 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
Jul 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
Jul 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
Jul 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
Jul 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
malfet
added a commit
to pytorch/torchchat
that referenced
this pull request
Jul 17, 2024
By patching runtime/jit.py with the contents of triton-lang/triton#3564 Should fix intermittent ``` SyntaxError: unterminated string literal (detected at line 1) ``` for good TODO: Delete `scripts/patch_trition.py` as well as changes to `install_requirements.sh` for good
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Under certain conditions where we apply another decorator on top of triton.jit, there is a small chance that in the decorator body there's a "def" substring in it. This is problematic when JITFunction parses function def using str.find('def'). This causes an issue compiling some pytorch model when we happened to decorate a jit function with a hash string argument which has "def" in it.
This diff tries a bit harder to locate the real "def" position using a more complex regex which should further reduce the chance to hit this error.