[Mlir-commits] [mlir] a677a17 - [mlir][py] Enable AsmState overload for operation.

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2023-09-25T12:25:08-07:00
New Revision: a677a173273f7735a3a5447a47851591b984e7aa

URL: https://github.com/llvm/llvm-project/commit/a677a173273f7735a3a5447a47851591b984e7aa
DIFF: https://github.com/llvm/llvm-project/commit/a677a173273f7735a3a5447a47851591b984e7aa.diff

LOG: [mlir][py] Enable AsmState overload for operation.

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/Bindings/Python/IRModule.h
    mlir/test/python/ir/value.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index c74b37a51c0df06..aad74f511e7ef11 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3479,6 +3479,8 @@ void mlir::python::populateIRCore(py::module &m) {
 
   py::class_<PyAsmState>(m, "AsmState", py::module_local())
       .def(py::init<PyValue &, bool>(), py::arg("value"),
+           py::arg("use_local_scope") = false)
+      .def(py::init<PyOperationBase &, bool>(), py::arg("op"),
            py::arg("use_local_scope") = false);
 
   //----------------------------------------------------------------------------

diff  --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 23338f7fdb38add..3ca7dd851961a46 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -759,6 +759,16 @@ class PyAsmState {
       mlirOpPrintingFlagsUseLocalScope(flags);
     state = mlirAsmStateCreateForValue(value, flags);
   }
+
+  PyAsmState(PyOperationBase &operation, bool useLocalScope) {
+    flags = mlirOpPrintingFlagsCreate();
+    // The OpPrintingFlags are not exposed Python side, create locally and
+    // associate lifetime with the state.
+    if (useLocalScope)
+      mlirOpPrintingFlagsUseLocalScope(flags);
+    state =
+        mlirAsmStateCreateForOperation(operation.getOperation().get(), flags);
+  }
   ~PyAsmState() {
     mlirOpPrintingFlagsDestroy(flags);
   }

diff  --git a/mlir/test/python/ir/value.py b/mlir/test/python/ir/value.py
index 2a47c8d820eaf4f..ddf653dcce27804 100644
--- a/mlir/test/python/ir/value.py
+++ b/mlir/test/python/ir/value.py
@@ -165,8 +165,8 @@ def testValuePrintAsOperand():
             # CHECK: Value(%[[VAL2:.*]] = "custom.op2"() : () -> i32)
             print(value2)
 
-            f = func.FuncOp("test", ([i32, i32], []))
-            entry_block1 = Block.create_at_start(f.operation.regions[0], [i32, i32])
+            topFn = func.FuncOp("test", ([i32, i32], []))
+            entry_block1 = Block.create_at_start(topFn.operation.regions[0], [i32, i32])
 
             with InsertionPoint(entry_block1):
                 value3 = Operation.create("custom.op3", results=[i32]).results[0]
@@ -201,7 +201,7 @@ def testValuePrintAsOperand():
 
         print("With AsmState")
         # CHECK-LABEL: With AsmState
-        state = AsmState(value3, use_local_scope=True)
+        state = AsmState(topFn.operation, use_local_scope=True)
         # CHECK: %0
         print(value3.get_name(state=state))
         # CHECK: %1