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.
|
|
||
| # TODO: handle these - need to build CompilerOptions | ||
|
|
||
| codegen: str = "loopy" |
| 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'
| ContextClass = LoopyCodegenContext | ||
| elif compiler_parameters.codegen == "mlir": | ||
| raise NotImplementedError("Class is still being implemented.") | ||
|
|
||
| context = ContextClass( |
There was a problem hiding this comment.
| ContextClass = LoopyCodegenContext | |
| elif compiler_parameters.codegen == "mlir": | |
| raise NotImplementedError("Class is still being implemented.") | |
| context = ContextClass( | |
| make_context = LoopyCodegenContext | |
| elif compiler_parameters.codegen == "mlir": | |
| raise NotImplementedError("Class is still being implemented.") | |
| context = make_context( |
| 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.
|
|
||
| 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
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.