From 61fa5c2c5fe1b295ab4aa8dec7d065bbf847a1c9 Mon Sep 17 00:00:00 2001 From: PragmaTwice Date: Fri, 12 Sep 2025 12:07:10 +0800 Subject: [PATCH 1/7] [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 ""("{{.*}}[[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 ""("{{.*}}[[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..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 ""("{{.*}}[[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..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 ""("{{.*}}[[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 Date: Fri, 12 Sep 2025 13:00:40 +0800 Subject: [PATCH 2/7] 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 Date: Fri, 12 Sep 2025 13:04:03 +0800 Subject: [PATCH 3/7] 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 Date: Fri, 12 Sep 2025 13:53:18 +0800 Subject: [PATCH 4/7] 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 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); From 6cf67d39ecc105df874e026a46291eb571dee346 Mon Sep 17 00:00:00 2001 From: Twice Date: Fri, 12 Sep 2025 14:07:03 +0800 Subject: [PATCH 5/7] Update mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp Co-authored-by: Maksim Levental --- mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp index 478c8eb2d1909..e69ff0cd7bf27 100644 --- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp +++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp @@ -1049,7 +1049,7 @@ static std::string makeDocStringForOp(const Operator &op) { std::string docString = "\n"; llvm::raw_string_ostream os(docString); raw_indented_ostream identedOs(os); - os << R"( """)" << "\n"; + os << R"( r""")" << "\n"; identedOs.printReindented(desc, " "); if (!StringRef(desc).ends_with("\n")) os << "\n"; From 25d1ae261aca4312d4d1008007b39f74d78e0366 Mon Sep 17 00:00:00 2001 From: PragmaTwice Date: Fri, 12 Sep 2025 14:08:40 +0800 Subject: [PATCH 6/7] fix test case --- mlir/test/mlir-tblgen/op-python-bindings.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/test/mlir-tblgen/op-python-bindings.td b/mlir/test/mlir-tblgen/op-python-bindings.td index 9f36d2b1ad0f5..a78ec09c531ba 100644 --- a/mlir/test/mlir-tblgen/op-python-bindings.td +++ b/mlir/test/mlir-tblgen/op-python-bindings.td @@ -254,7 +254,7 @@ def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [Fir // CHECK: class DescriptionOp(_ods_ir.OpView): -// CHECK: """ +// CHECK: r""" // CHECK: This is a long description. // CHECK: It has multiple lines. // CHECK: A code block (to test the indent). From 8b304c201fa40194cf34faf3d586632e14745277 Mon Sep 17 00:00:00 2001 From: PragmaTwice Date: Tue, 16 Sep 2025 10:52:20 +0800 Subject: [PATCH 7/7] fix regex --- mlir/test/python/ir/auto_location.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/test/python/ir/auto_location.py b/mlir/test/python/ir/auto_location.py index ed2f255af2259..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":{{\d+}}: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 ""("{{.*}}[[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 ""("{{.*}}[[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":{{\d+}}:4 to :235) at callsite("testInferLocations..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 ""("{{.*}}[[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..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 ""("{{.*}}[[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":{{\d+}}:4 to :235)) + # CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":{{[0-9]+}}:4 to :235)) # fmt: on foo()