After talking a bit with @jdd, we are starting to consider a Python-based system for defining ops and performing some relatively simple rewrites. It seems like this is a potential place to leverage PDL. Before diving in, we thought it would be prudent to ask the community: is anyone else interested in using Python as a frontend to PDL? If we do dive into this, we would certainly hope to collaborate.
To give a little more background on what we’re trying to achieve, it’s somewhat similar to the use case of defining a bunch of named linalg ops. Except in our case, it is a hardware-oriented high-level dialect. We hope to use Python both to define the high-level model, as well as how it rewrites into a lower-level dialect. Sorry if that is a bit hand-wavy, but this is very nascent.
I’m interested in seeing how a frontend for PDL would look like!
I’m also involved in another project about defining dialects dynamically (potentially from Python), which seems like another aspect you’re interested in?
Thanks for posting this. I’ve recently been planning how we are going to target PDL from various different angles. There are two main ones that I am particularly interested in: custom and python. A custom frontend has various benefits, namely that we can architect it in such a way that it can replace the current of use of tablegen for DRR (whereas a python frontend is something that wouldn’t work as a tenable replacement). I would happy to sync up and collaborate on developing something upstream that can work both for your usecases and the ones that I have in mind.
+100. I’m going to be doing something that sounds very similar in npcomp soon for the problem of lowering torch kernel calls (“aten ops”) to linalg named ops on tensors (possibly custom ones just for the for that kernel call if it does something unique).
To use hardware design terminology, the “design entry” Python API we’re envisioning allows designers to define some “design elements” with only the interface (modules and port lists, in hw jargon). We’re thinking that these could be MLIR operations with a little hardware-specific sugar. The designer would then be able choose a particular implementation, which we’re thinking would map directly to an MLIR lowering. Example: all IEEE 754 floating point units have the same interface – it follows from the IEEE spec, so that’d be the “design element”. There are a bunch of different ways to implement FP hardware with a series of trade-offs, each of which would be an “implementation” (MLIR lowering).
The designer would have access to all C++ MLIR/CIRCT operations/attributes/types to use as a standard library.
I’ve written a sketch of what we’d like to do, which may or may not be valid Python and/or correct.
@circt.module
class PolynomialCompute:
"""Module to compute ax^3 + bx^2 + cx + d for design-time coefficients."""
taps = Output(List[int])
y = Output(int)
def __init__(self, coefficients: List[int])
"""'coefficients' is in 'd' -> 'a' order."""
self.__coefficients = coefficients
@implementation("yuckyLegacyDevices")
def old_style(self, x: int):
"""Implement this module for input 'x'."""
# Output a 'tap' for each partial sum.
runningPower: List[int] = List()
for coeff in self.__coefficients:
if len(self.taps) == 0:
partialSum = 0
currPow = 1
else
partialSum = taps[-1]
currPow = circt.math.mult(x, runningPower[-1])
runningPower.append(currPow)
newPartialSum = circt.math.add(
partialSum,
circt.math.mult(coeff, currPow))
self.taps.append(newPartialSum)
# Final output
self.y = self.taps[-1]
@implementation("special_fpga_with_new_hotness")
def new_hotness(self, x: int):
"""Implement polynomial computes for new FPGA device."""
for i, coeff in enumerate(self.__coefficients):
partialSum = 0 if len(self.taps) > 0 else taps[-1]
xi = SuperCoolPowerCompute(i).create(x)
newPartialSum = circt.math.add(
partialSum,
circt.math.mult(coeff, xi.y))
self.taps.append(newPartialSum)
# Final output
self.y = self.taps[-1]
@circt.externmodule
class SuperCoolPowerCompute:
"""Super cool new dedicated hardware to compute exponents, but only if the
exponent is known at compile time."""
x = Input(int)
y = Output(int)
mode = Parameter(str) # Parameters would map to MLIR attributes
exp = Parameter(int)
def __init__(self, exp: int):
self.exp = exp
self.mode = "exp"
I wrote this example without any regard to PDL since I also don’t know what that would look like or if it would have any benefit here.
Using PDL to enable lightweight rewriting in the host language. I was thinking about this as well, and it seems pretty straightforward given that PDL is just another dialect.
“Custom” operations in the host language. @prithayan implemented the generator callout pass which currently is designed to call an external tool, but I think could easily be adapted for use with an in-process generator. I’m not sure this would need to rely on PDL, as opposed to some simpler mechanism, like just registering a generator for a particular schema.
I’m currently working on a project that has similar goals in mind, and would love to collaborate!
One part of the project is to provide an ad-hoc way to register custom operations/types/dialects at runtime. This could be used from Python to declare new operations for instance, which I think would interest you. We also would like to have a nice interaction for that in Python, so we can write “my_op.my_field” to directly access the corresponding field of the MLIR operation.
Another thing we would like to do is to provide ways of defining transformations (and rewrite rules) in Python. This could either be done by providing a PDL frontend in Python, or by providing a generic Python API so we could write more complex transformations from Python. For instance, this is how we would create a constant operation:
The main idea would be to define transformations in Python using this API, and then register those transformations in MLIR (still in Python). This could then be used to create more specialized frontends like the one you want.
Thanks all for weighing in. It’s encouraging to see others are interested in this.
After thinking about it some more, I want to be clear that there are indeed two things at play: on the one hand, modelling some high-level operations in Python, and on the other, rewriting/lowering/generating low-level operations from them.
We are still very much interested in the high-level modelling opportunities with Python, and that is going to be the immediate focus.
Regarding how to rewrite/lower/generate these into low-level operations, we are starting to think we may not need the full power of PDL, yet. We already have a decent Python API for constructing the low-level operations we care about, and these rewrites are going to be simple enough that we can probably just call the low-level operation builders from the higher-level operations.
In other words, I think we are trending towards something more like the opdsl for Linalg rather than a full-blown PDL-based system. Down the road, I can absolutely envision using PDL to enable some really cool transformations, but for the immediate use-case, we are going to see how far we can go without building out a proper Python frontend for PDL. Since it seems like there is a lot of similar work going on here, hopefully we can share what works well and what doesn’t, and still collaborate on whatever would be useful upstream.
