def PrintOpStats : Pass<"print-op-stats"> {
let summary = "Print statistics of operations";
let constructor = "mlir::createPrintOpStatsPass()";
+ let options = [
+ Option<"printAsJSON", "json", "bool", /*default=*/"false",
+ "print the stats as JSON">
+ ];
}
def SCCP : Pass<"sccp"> {
// Print summary of op stats.
void printSummary();
+ // Print symmary of op stats in JSON.
+ void printSummaryInJSON();
+
private:
llvm::StringMap<int64_t> opCount;
raw_ostream &os;
opCount.clear();
// Compute the operation statistics for the currently visited operation.
- getOperation()->walk([&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
- printSummary();
+ getOperation()->walk(
+ [&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
+ if (printAsJSON) {
+ printSummaryInJSON();
+ } else
+ printSummary();
}
void PrintOpStatsPass::printSummary() {
}
}
+void PrintOpStatsPass::printSummaryInJSON() {
+ SmallVector<StringRef, 64> sorted(opCount.keys());
+ llvm::sort(sorted);
+
+ os << "{\n";
+
+ for (unsigned i = 0, e = sorted.size(); i != e; ++i) {
+ const auto &key = sorted[i];
+ os << " \"" << key << "\" : " << opCount[key];
+ if (i != e - 1)
+ os << ",\n";
+ else
+ os << "\n";
+ }
+ os << "}\n";
+}
+
std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os) {
return std::make_unique<PrintOpStatsPass>(os);
}
mlirOpPassManagerAddOwnedPass(nestedFuncPm, printOpStatPass);
// Print the top level pass manager
- // CHECK: Top-level: builtin.module(func.func(print-op-stats))
+ // CHECK: Top-level: builtin.module(func.func(print-op-stats{json=false}))
fprintf(stderr, "Top-level: ");
mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
NULL);
fprintf(stderr, "\n");
// Print the pipeline nested one level down
- // CHECK: Nested Module: func.func(print-op-stats)
+ // CHECK: Nested Module: func.func(print-op-stats{json=false})
fprintf(stderr, "Nested Module: ");
mlirPrintPassPipeline(nestedModulePm, printToStderr, NULL);
fprintf(stderr, "\n");
// Try parse a pipeline.
MlirLogicalResult status = mlirParsePassPipeline(
mlirPassManagerGetAsOpPassManager(pm),
- mlirStringRefCreateFromCString("builtin.module(func.func(print-op-stats),"
- " func.func(print-op-stats))"));
+ mlirStringRefCreateFromCString(
+ "builtin.module(func.func(print-op-stats{json=false}),"
+ " func.func(print-op-stats{json=false}))"));
// Expect a failure, we haven't registered the print-op-stats pass yet.
if (mlirLogicalResultIsSuccess(status)) {
fprintf(
mlirRegisterTransformsPrintOpStats();
status = mlirParsePassPipeline(
mlirPassManagerGetAsOpPassManager(pm),
- mlirStringRefCreateFromCString("builtin.module(func.func(print-op-stats),"
- " func.func(print-op-stats))"));
+ mlirStringRefCreateFromCString(
+ "builtin.module(func.func(print-op-stats{json=false}),"
+ " func.func(print-op-stats{json=false}))"));
// Expect a failure, we haven't registered the print-op-stats pass yet.
if (mlirLogicalResultIsFailure(status)) {
fprintf(stderr,
exit(EXIT_FAILURE);
}
- // CHECK: Round-trip: builtin.module(func.func(print-op-stats),
- // func.func(print-op-stats))
+ // CHECK: Round-trip: builtin.module(func.func(print-op-stats{json=false}),
+ // func.func(print-op-stats{json=false}))
fprintf(stderr, "Round-trip: ");
mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
NULL);
--- /dev/null
+// RUN: mlir-opt -allow-unregistered-dialect -print-op-stats=json %s -o=/dev/null 2>&1 | FileCheck %s
+
+func.func @main(tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> {
+^bb0(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>):
+ %0 = arith.addf %arg0, %arg1 : tensor<4xf32>
+ %1 = arith.addf %arg0, %arg1 : tensor<4xf32>
+ %2 = arith.addf %arg0, %arg1 : tensor<4xf32>
+ %3 = arith.addf %arg0, %arg1 : tensor<4xf32>
+ %4 = arith.addf %arg0, %arg1 : tensor<4xf32>
+ %5 = arith.addf %arg0, %arg1 : tensor<4xf32>
+ %10 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %11 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %12 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %13 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %14 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %15 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %16 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %17 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %18 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %19 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %20 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %21 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %22 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %23 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %24 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %25 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %26 = "xla.add"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ %30 = "long_op_name"(%0, %arg1) : (tensor<4xf32>,tensor<4xf32>)-> tensor<4xf32>
+ return %1 : tensor<4xf32>
+}
+
+// CHECK: {
+// CHECK: "arith.addf" : 6,
+// CHECK: "func.return" : 1,
+// CHECK: "long_op_name" : 1,
+// CHECK: "xla.add" : 17
+// CHECK: }
# A first import is expected to fail because the pass isn't registered
# until we import mlir.transforms
try:
- pm = PassManager.parse("builtin.module(func.func(print-op-stats))")
+ pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
# TODO: this error should be propagate to Python but the C API does not help right now.
# CHECK: error: 'print-op-stats' does not refer to a registered pass or pass pipeline
except ValueError as e:
- # CHECK: ValueError exception: invalid pass pipeline 'builtin.module(func.func(print-op-stats))'.
+ # CHECK: ValueError exception: invalid pass pipeline 'builtin.module(func.func(print-op-stats{json=false}))'.
log("ValueError exception:", e)
else:
log("Exception not produced")
# This will register the pass and round-trip should be possible now.
import mlir.transforms
- pm = PassManager.parse("builtin.module(func.func(print-op-stats))")
- # CHECK: Roundtrip: builtin.module(func.func(print-op-stats))
+ pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
+ # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))
log("Roundtrip: ", pm)
run(testParseSuccess)
# CHECK-LABEL: TEST: testRun
def testRunPipeline():
with Context():
- pm = PassManager.parse("print-op-stats")
+ pm = PassManager.parse("print-op-stats{json=false}")
module = Module.parse(r"""func.func @successfulParse() { return }""")
pm.run(module)
# CHECK: Operations encountered: