Allow customization of the subscript operator for triton values #7239
Conversation
|
@Mogball are you the right person to look at this? |
Currently, it invokes the right method but it doesn't let you customize it with the |
Makes sense. Does this mean that other binop methods will need to go through the same |
|
Thanks for the frontend contributions by the way. They're much appreciated |
That's what I was thinking, yes. |
| return lhs[slices] | ||
|
|
||
| def visit_Subscript_Store(self, node, value): | ||
| assert isinstance(node.ctx, ast.Store) | ||
| lhs = self.visit(node.value) | ||
| slices = self.visit(node.slice) | ||
| assert isinstance(lhs, language.tuple) | ||
| lhs.__setitem__(slices, value) | ||
| fn = BoundJITMethod(lhs, lhs.__setitem__) if isinstance(lhs.__setitem__, JITFunction) else lhs.__setitem__ |
There was a problem hiding this comment.
I'm a little frustrated by this. I would have thought that __setitem__ would be a bound method containing a JITFunction however this doesn't appear to be the case (it's just a JITFunction) so I need to convert it into a BoundJITMethod in order to call it. I'm not sure if this is an issue cause by the _argegate wrapper on this is just an issue with python. Do you have any ideas?
There was a problem hiding this comment.
The method unfortunately can't be bound to the object until Triton parses the code, because it needs to be bound to the MLIR handle. I had changed visit_Attribute to emit obj.method where method is a JITFunction to bind obj into it, which is different than Python which binds upon object construction.
The methods have to be bound during visit_Attribute because Triton values are not reference semantic. This does mean that something like
f = obj.method
obj.x = 1234
f()Would not work as one would expect, if it was even valid code to begin with.
You can probably wrap this in a call_Method helper since it could end up being quite common.
There was a problem hiding this comment.
Makes sense thanks for explaining. I wonder if we can get closer to python by making @jit return a function rather than an object. I'll play around and make a follow up draft PR if that idea goes anywhere. For now I'll make a call_Method helper.
There was a problem hiding this comment.
Yeah that idea does work. I'll make a follow up draft PR once this lands.
| extra_kwargs = dict() | ||
| sig = inspect.signature(fn) | ||
| if '_semantic' in sig.parameters: | ||
| extra_kwargs["_semantic"] = self.semantic | ||
| if '_generator' in sig.parameters: | ||
| extra_kwargs['_generator'] = self |
There was a problem hiding this comment.
Now that I've made both _semantic and _generator optional. Do we need the if statement at all? I guess the only differences are:
- how the results are wrapped (although this probably should be made consistent)
- whether the error is wrapped in CompilationError (again this should probably be made consistent)
- Theoretically the user might have used
_semanticand_generatoras arguments without intending them to be overwritten by us which could lead to weird errors. Although if we were really concerned about this I think the fix would be to mangle the names even more.
This allows users to provide custom setitem and getitem functions in order to override the subscript operator for their triton classes.
New contributor declaration
I am not making a trivial change, such as fixing a typo in a comment.
I have written a PR description following these
rules.
I have run
pre-commit run --from-ref origin/main --to-ref HEAD.Select one of the following.
/testforlittests/unittestfor C++ tests/python/testfor end-to-end testsFILL THIS IN.Select one of the following.
littests.littests I have added follow these best practices,including the "tests should be minimal" section. (Usually running Python code
and using the instructions it generates is not minimal.)