[Mlir-commits] [mlir] [MLIR][Python] Add docstring for generated python op classes (PR #158198)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 11 22:53:49 PDT 2025
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/158198
>From 61fa5c2c5fe1b295ab4aa8dec7d065bbf847a1c9 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Fri, 12 Sep 2025 12:07:10 +0800
Subject: [PATCH 1/4] [MLIR][Python] Add docstring for generated python op
classes
---
mlir/test/python/ir/auto_location.py | 6 ++--
mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 30 +++++++++++++++++--
2 files changed, 30 insertions(+), 6 deletions(-)
diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py
index 01b5542119b4e..24b4fb076afe0 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":397: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":649: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":397: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":649: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":397:4 to :235))
+ # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":649:4 to :235))
# fmt: on
foo()
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 6a7aa9e3432d5..e45a142cf9d38 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"
@@ -62,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";
@@ -1034,9 +1036,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 = llvm::join(llvm::split(desc, R"(""")"), R"(\"\"\")");
+
+ std::string docString = "\n";
+ llvm::raw_string_ostream os(docString);
+ raw_indented_ostream identedOs(os);
+ os << 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) {
>From 0e199e4a35144ce0b02bcd7859de314f102c1d46 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Fri, 12 Sep 2025 13:00:40 +0800
Subject: [PATCH 2/4] Add a tblgen test case
---
mlir/test/mlir-tblgen/op-python-bindings.td | 26 +++++++++++++++++++++
1 file changed, 26 insertions(+)
diff --git a/mlir/test/mlir-tblgen/op-python-bindings.td b/mlir/test/mlir-tblgen/op-python-bindings.td
index 3ec69c33b4bb9..5775e5c0eeda9 100644
--- a/mlir/test/mlir-tblgen/op-python-bindings.td
+++ b/mlir/test/mlir-tblgen/op-python-bindings.td
@@ -252,6 +252,32 @@ def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [Fir
// CHECK: def derive_result_types_variadic_op(res, _gen_res_1, type_, *, loc=None, ip=None)
// CHECK: return _get_op_result_or_op_results(DeriveResultTypesVariadicOp(res=res, _gen_res_1=_gen_res_1, type_=type_, loc=loc, ip=ip))
+
+// CHECK: class DescriptionOp(_ods_ir.OpView):
+// CHECK: """
+// 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: """
+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
+ }
+ ```
+ }];
+}
+
// CHECK: @_ods_cext.register_operation(_Dialect)
// CHECK: class EmptyOp(_ods_ir.OpView):
// CHECK-LABEL: OPERATION_NAME = "test.empty"
>From 7795d9ea4ab0f95e57f36836eda69f7f2bb1f4b5 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Fri, 12 Sep 2025 13:04:03 +0800
Subject: [PATCH 3/4] refine the test case
---
mlir/test/mlir-tblgen/op-python-bindings.td | 2 ++
1 file changed, 2 insertions(+)
diff --git a/mlir/test/mlir-tblgen/op-python-bindings.td b/mlir/test/mlir-tblgen/op-python-bindings.td
index 5775e5c0eeda9..9f36d2b1ad0f5 100644
--- a/mlir/test/mlir-tblgen/op-python-bindings.td
+++ b/mlir/test/mlir-tblgen/op-python-bindings.td
@@ -263,6 +263,7 @@ def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [Fir
// CHECK: test.yield
// CHECK: }
// CHECK: ```
+// CHECK: Add \"\"\" will not terminate the description.
// CHECK: """
def DescriptionOp : TestOp<"description"> {
let description = [{
@@ -275,6 +276,7 @@ def DescriptionOp : TestOp<"description"> {
test.yield
}
```
+ Add """ will not terminate the description.
}];
}
>From 1cc3d1441438a5b681bab5b39a9f87e1a2358281 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Fri, 12 Sep 2025 13:53:18 +0800
Subject: [PATCH 4/4] fix replace
---
mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index e45a142cf9d38..478c8eb2d1909 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -21,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;
@@ -1043,7 +1044,7 @@ static std::string makeDocStringForOp(const Operator &op) {
auto desc = op.getDescription().rtrim(" \t").str();
// Replace all """ with \"\"\" to avoid early termination of the literal.
- desc = llvm::join(llvm::split(desc, R"(""")"), R"(\"\"\")");
+ desc = std::regex_replace(desc, std::regex(R"(""")"), R"(\"\"\")");
std::string docString = "\n";
llvm::raw_string_ostream os(docString);
More information about the Mlir-commits
mailing list