[Mlir-commits] [mlir] 10d0d95 - [MLIR][Python] Add docstring for generated python op classes (#158198)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 15 20:16:32 PDT 2025
Author: Twice
Date: 2025-09-16T11:16:28+08:00
New Revision: 10d0d955e22be43abab8fe7e339e99068aaac030
URL: https://github.com/llvm/llvm-project/commit/10d0d955e22be43abab8fe7e339e99068aaac030
DIFF: https://github.com/llvm/llvm-project/commit/10d0d955e22be43abab8fe7e339e99068aaac030.diff
LOG: [MLIR][Python] Add docstring for generated python op classes (#158198)
This PR adds support in mlir-tblgen for generating docstrings for each
Python class corresponding to an MLIR op. The docstrings are currently
derived from the op’s description in ODS, with indentation adjusted to
display nicely in Python. This makes it easier for Python users to see
the op descriptions directly in their IDE or LSP while coding.
In the future, we can extend the docstrings to include explanations for
each method, attribute, and so on.
This idea was previously discussed in the `#mlir-python` channel on
Discord with @makslevental and @superbobry.
---------
Co-authored-by: Maksim Levental <maksim.levental at gmail.com>
Added:
Modified:
mlir/test/mlir-tblgen/op-python-bindings.td
mlir/test/python/ir/auto_location.py
mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
Removed:
################################################################################
diff --git a/mlir/test/mlir-tblgen/op-python-bindings.td b/mlir/test/mlir-tblgen/op-python-bindings.td
index 90feec9ed8d6b..81a4b2bbd1631 100644
--- a/mlir/test/mlir-tblgen/op-python-bindings.td
+++ b/mlir/test/mlir-tblgen/op-python-bindings.td
@@ -254,6 +254,34 @@ def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [Fir
// CHECK: op = DeriveResultTypesVariadicOp(res=res, _gen_res_1=_gen_res_1, type_=type_, loc=loc, ip=ip); results = op.results
// CHECK: return results if len(results) > 1 else (results[0] if len(results) == 1 else op)
+
+// CHECK: class DescriptionOp(_ods_ir.OpView):
+// CHECK: r"""
+// CHECK: This is a long description.
+// CHECK: It has multiple lines.
+// CHECK: A code block (to test the indent).
+// CHECK: ```mlir
+// CHECK: test.loop {
+// CHECK: test.yield
+// CHECK: }
+// CHECK: ```
+// CHECK: Add \"\"\" will not terminate the description.
+// CHECK: """
+def DescriptionOp : TestOp<"description"> {
+ let description = [{
+ This is a long description.
+ It has multiple lines.
+
+ A code block (to test the indent).
+ ```mlir
+ test.loop {
+ test.yield
+ }
+ ```
+ Add """ will not terminate the description.
+ }];
+}
+
// CHECK: @_ods_cext.register_operation(_Dialect)
// CHECK: class EmptyOp(_ods_ir.OpView):
// CHECK-LABEL: OPERATION_NAME = "test.empty"
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index a063aa972cc48..83168901341a8 100644
--- a/mlir/test/python/ir/auto_location.py
+++ b/mlir/test/python/ir/auto_location.py
@@ -51,7 +51,7 @@ def testInferLocations():
_cext.globals.register_traceback_file_inclusion(_arith_ops_gen.__file__)
three = arith.constant(IndexType.get(), 3)
# fmt: off
- # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":396:4 to :235) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":52:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))
+ # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":52:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))
# fmt: on
print(three.location)
@@ -60,14 +60,14 @@ def foo():
print(four.location)
# fmt: off
- # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":396:4 to :235) at callsite("testInferLocations.<locals>.foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":65:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))))
+ # CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235) at callsite("testInferLocations.<locals>.foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":65:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))))
# fmt: on
foo()
_cext.globals.register_traceback_file_exclusion(__file__)
# fmt: off
- # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":396:4 to :235))
+ # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235))
# fmt: on
foo()
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 21f712e85e6c0..169d550f44850 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -13,6 +13,7 @@
#include "OpGenHelpers.h"
+#include "mlir/Support/IndentedOstream.h"
#include "mlir/TableGen/GenInfo.h"
#include "mlir/TableGen/Operator.h"
#include "llvm/ADT/StringSet.h"
@@ -20,6 +21,7 @@
#include "llvm/Support/FormatVariadic.h"
#include "llvm/TableGen/Error.h"
#include "llvm/TableGen/Record.h"
+#include <regex>
using namespace mlir;
using namespace mlir::tblgen;
@@ -61,10 +63,11 @@ from ._{0}_ops_gen import _Dialect
/// Template for operation class:
/// {0} is the Python class name;
-/// {1} is the operation name.
+/// {1} is the operation name;
+/// {2} is the docstring for this operation.
constexpr const char *opClassTemplate = R"Py(
@_ods_cext.register_operation(_Dialect)
-class {0}(_ods_ir.OpView):
+class {0}(_ods_ir.OpView):{2}
OPERATION_NAME = "{1}"
)Py";
@@ -1031,9 +1034,31 @@ static void emitValueBuilder(const Operator &op,
}
}
+/// Retrieve the description of the given op and generate a docstring for it.
+static std::string makeDocStringForOp(const Operator &op) {
+ if (!op.hasDescription())
+ return "";
+
+ auto desc = op.getDescription().rtrim(" \t").str();
+ // Replace all """ with \"\"\" to avoid early termination of the literal.
+ desc = std::regex_replace(desc, std::regex(R"(""")"), R"(\"\"\")");
+
+ std::string docString = "\n";
+ llvm::raw_string_ostream os(docString);
+ raw_indented_ostream identedOs(os);
+ os << R"( r""")" << "\n";
+ identedOs.printReindented(desc, " ");
+ if (!StringRef(desc).ends_with("\n"))
+ os << "\n";
+ os << R"( """)" << "\n";
+
+ return docString;
+}
+
/// Emits bindings for a specific Op to the given output stream.
static void emitOpBindings(const Operator &op, raw_ostream &os) {
- os << formatv(opClassTemplate, op.getCppClassName(), op.getOperationName());
+ os << formatv(opClassTemplate, op.getCppClassName(), op.getOperationName(),
+ makeDocStringForOp(op));
// Sized segments.
if (op.getTrait(attrSizedTraitForKind("operand")) != nullptr) {
More information about the Mlir-commits
mailing list