PR: Refactoring PyOP3 Codegen - #5392
Conversation
- Goal: Create an interface to a code generation context (MLIR or Loopy) - Status: Interface created, battling PETSc bug before cleaning more. - Goal: Integrate MLIR for auto-generation - Status: Was working but transitioning to pyop3->mlir pipeline as opposed to pym->mlir. Refactoring process is ongoing.
requirements tracking pyop3 && typing hints
Merged Connor's update with new setup. Current issue with values loaded into the buffers. Occurs even while using old lower/loopy.py Debugging in process.
PyOP3 updates and path traversals reflected in the new lowering structure
Still need to integrate dtype support for all op3 expressions. Removed abstractmethod while testing.
connorjward
left a comment
connorjward
left a comment
There was a problem hiding this comment.
Key point on an initial first pass is I would like to limit the scope of what the context objects know about.
| for e in pyop3.collections.as_tuple(ex): # TODO: get rid of this loop | ||
| # context manager? | ||
| context.set_temporary_shapes(_collect_temporary_shapes(e)) | ||
| _compile(e, loop_indices, context) |
There was a problem hiding this comment.
This code is really disgusting but not your problem!
| else: | ||
| cs_expr = (insn,) | ||
|
|
||
| if compiler_parameters.codegen == "loopy": |
There was a problem hiding this comment.
I think I prefer 'backend' to 'codegen'
There was a problem hiding this comment.
That's fine, I will switch. Was only wary that backend is used often but agree that it makes more sense.
There was a problem hiding this comment.
Terminology is really confusing around here. We have:
- compile
- codegen
- lower
which are all in some sense interchangeable. I really want to use "compile" as the single term for everything but this conflicts with the other sense of "compiler" (i.e. GCC etc).
I think in my main branch I will do the following renaming:
pyop3/lowertopyop3/compilepyop3/compile.pytopyop3/cc.py
File naming isn't important for this PR, but I think I will make that change soon and you may hit git conflicts.
| return ctx | ||
|
|
||
|
|
||
| # NOTE: Not a big fan of how compile sits in this file. |
There was a problem hiding this comment.
Yeah I would expect to put it in codegen.py. This file can just contain the abstract class.
There was a problem hiding this comment.
Moved _compile to codegen.py.
Only point of contention - and why it was not already there - is that parse_loop_properly_this_time uses _compile in the context.py file. As such, _compile now has to be imported. This must be a local import, to avoid a circular import.
It feels like a code smell but it does work.
Possible steps for resolving the potential smell would be:
- Move
_compilefunction to another separate file (I think this would bloat the directory) - Rewrite
parse_loop_properly_this_time
Else, we can leave it for later, it works fine. May just be a micro-optimisation.
There was a problem hiding this comment.
parse_loop_properly_this_time (which I think I would like you to rename) is another function that I don't think should belong to the context and ideally could be rewritten in a backend agnostic fashion.
|
|
||
| return indices | ||
|
|
||
| def compile_standalone_function( |
There was a problem hiding this comment.
When I designed the 'context' class I imagined that it would be entirely independent from any of the pyop3 'language'. It would only see the buffers. For example here I think we could amend the add_function_call to take in more arguments. The visitor for when we hit standalone function can then be
args = [(a.buffer_view, spec.intent) for a, spec in zip(call.arguments, call.argspec, strict=True)]
context.add_function_call(call.function.code, args)The crucial points are:
- The context never sees the
StandaloneCalledFunctiontype - The visitor for
StandaloneCalledFunctioncan remain generic for different backends
| subkernel = call.function.code.with_entrypoints(frozenset()) | ||
| self.add_subkernel(subkernel) | ||
|
|
||
| def compile_petsc_mat( |
There was a problem hiding this comment.
Same, I'd keep it generic. You can just raise NotImplementedError for non-loopy
- rename compiler option `codegen` -> `backend` - documenting `backend` compiler option - functional variable renaming
- moving dispatch _compile function from context.py -> codegen.py - local import in _parse_loop to avoid circular import - moving dispatch _collect_temporary_shapes from context.py -> transform.py
2 participants
Refactoring
pyop3/lower/PR presents abstraction in the
lower/design to allow for alternative code generation backends.Principally,
loopy.pyhas been broken into three separate files:codegen.py- context orchestrator which returns lowered IR back to translation layercontext.py- generic class and high-level traversal of axis trees to generate pyop3 expressions.(possible that class and traversal has too many responsibilities).
loopy.py-implements backend-specific lowering from pyop3 to respective target IRsThis PR serves as a stepping stone to the introduction of MLIR as a code generation backend.
A new
mlir.pymodule will subclasscontext.pythat will construct MLIR from PyOP3 expressions, using xDSL.