Skip to content

Commit 00a93e6

Browse files
authored
[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.
1 parent c271ba7 commit 00a93e6

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

mlir/lib/Bindings/Python/IRCore.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1534,7 +1534,7 @@ py::object PyOperation::create(const std::string &name,
15341534
PyOperation::createDetached(location->getContext(), operation);
15351535
maybeInsertOperation(created, maybeIp);
15361536

1537-
return created->createOpView();
1537+
return created.getObject();
15381538
}
15391539

15401540
py::object PyOperation::clone(const py::object &maybeIp) {

0 commit comments

Comments
 (0)