[Mlir-commits] [mlir] e0928ab - [mlir][c] Add mlirOperationPrintWithState

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2023-11-03T16:29:03-07:00
New Revision: e0928abb48e465b82fcd58f92c75bc9d1d9f64da

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

LOG: [mlir][c] Add mlirOperationPrintWithState

Enable passing in MlirAsmState optionally (allow for passing in null) to
allow using the more efficient print calling API. The existing print
behavior results in a new AsmState is implicitly created by walking the
parent op and renumbering values. This makes the cost more explicit and
avoidable (by reusing an AsmState).

Added: 
    

Modified: 
    mlir/include/mlir-c/IR.h
    mlir/lib/CAPI/IR/IR.cpp
    mlir/test/CAPI/ir.c

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 5659230a03d8ce3..413eaa6aa3fe0ec 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -667,6 +667,13 @@ MLIR_CAPI_EXPORTED void mlirOperationPrintWithFlags(MlirOperation op,
                                                     MlirStringCallback callback,
                                                     void *userData);
 
+/// Same as mlirOperationPrint but accepts AsmState controlling the printing
+/// behavior as well as caching computed names.
+MLIR_CAPI_EXPORTED void mlirOperationPrintWithState(MlirOperation op,
+                                                    MlirAsmState state,
+                                                    MlirStringCallback callback,
+                                                    void *userData);
+
 /// Same as mlirOperationPrint but writing the bytecode format.
 MLIR_CAPI_EXPORTED void mlirOperationWriteBytecode(MlirOperation op,
                                                    MlirStringCallback callback,

diff  --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 0a5151751873f2b..d1ee1b774c34478 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -678,6 +678,14 @@ void mlirOperationPrintWithFlags(MlirOperation op, MlirOpPrintingFlags flags,
   unwrap(op)->print(stream, *unwrap(flags));
 }
 
+void mlirOperationPrintWithState(MlirOperation op, MlirAsmState state,
+                                 MlirStringCallback callback, void *userData) {
+  detail::CallbackOstream stream(callback, userData);
+  if (state.ptr)
+    unwrap(op)->print(stream, *unwrap(state));
+  unwrap(op)->print(stream);
+}
+
 void mlirOperationWriteBytecode(MlirOperation op, MlirStringCallback callback,
                                 void *userData) {
   detail::CallbackOstream stream(callback, userData);

diff  --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index c6425f80a8bce9c..8d5dcbf62e85e2b 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -482,11 +482,14 @@ static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
   fprintf(stderr, "Op print with all flags: ");
   mlirOperationPrintWithFlags(operation, flags, printToStderr, NULL);
   fprintf(stderr, "\n");
+  fprintf(stderr, "Op print with state: ");
+  MlirAsmState state = mlirAsmStateCreateForOperation(parentOperation, flags);
+  mlirOperationPrintWithState(operation, state, printToStderr, NULL);
+  fprintf(stderr, "\n");
   // clang-format off
   // CHECK: Op print with all flags: %{{.*}} = "arith.constant"() <{value = 0 : index}> {elts = dense_resource<__elided__> : tensor<4xi32>} : () -> index loc(unknown)
   // clang-format on
 
-  MlirAsmState state = mlirAsmStateCreateForOperation(parentOperation, flags);
   fprintf(stderr, "With state: |");
   mlirValuePrintAsOperand(value, state, printToStderr, NULL);
   // CHECK: With state: |%0|