commit 00a93e6207ac053e7a42368a52cf51a744ecd139
author Peter Hawkins <phawkins@google.com> Mon Nov 04 20:19:18 2024 -0500
committer GitHub <noreply@github.com> Tue Nov 05 02:19:18 2024 +0100
tree 3b08b1f72d52ff143169eb182238203e844e713c
parent c271ba7f7961f7c2ba911853113e67ed6237a77b

[mlir:python] Change PyOperation::create to actually return a PyOperation. (#114542)

In the tablegen-generated Python bindings, we typically see a pattern
like:
```
class ConstantOp(_ods_ir.OpView):
  ...
  def __init__(self, value, *, loc=None, ip=None):
    ...
    super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))
```

i.e., the generated code calls `OpView.__init__()` with the output of
`build_generic`. The purpose of `OpView` is to wrap another operation
object, and `OpView.__init__` can accept any `PyOperationBase` subclass,
and presumably the intention is that `build_generic` returns a
`PyOperation`, so the user ends up with a `PyOpView` wrapping a
`PyOperation`.

However, `PyOpView::buildGeneric` calls `PyOperation::create`, which
does not just build a PyOperation, but it also calls `createOpView` to
wrap that operation in a subclass of `PyOpView` and returns that view.
But that's rather pointless: we called this code from the constructor of
an `OpView` subclass, so we already have a view object ready to go; we
don't need to build another one!

If we change `PyOperation::create` to return the underlying
`PyOperation`, rather than a view wrapper, we can save allocating a
useless `PyOpView` object for each ODS-generated Python object.

This saves approximately 1.5s of Python time in a JAX LLM benchmark that
generates a mixture of upstream dialects and StableHLO.