github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/third_party/mlir/lib/Analysis/OpStats.cpp (about)

     1  //===- OpStats.cpp - Prints stats of operations in module -----------------===//
     2  //
     3  // Copyright 2019 The MLIR Authors.
     4  //
     5  // Licensed under the Apache License, Version 2.0 (the "License");
     6  // you may not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //   http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  // =============================================================================
    17  
    18  #include "mlir/IR/Module.h"
    19  #include "mlir/IR/Operation.h"
    20  #include "mlir/IR/OperationSupport.h"
    21  #include "mlir/Pass/Pass.h"
    22  #include "llvm/ADT/DenseMap.h"
    23  #include "llvm/Support/Format.h"
    24  #include "llvm/Support/raw_ostream.h"
    25  
    26  using namespace mlir;
    27  
    28  namespace {
    29  struct PrintOpStatsPass : public ModulePass<PrintOpStatsPass> {
    30    explicit PrintOpStatsPass(llvm::raw_ostream &os = llvm::errs()) : os(os) {}
    31  
    32    // Prints the resultant operation statistics post iterating over the module.
    33    void runOnModule() override;
    34  
    35    // Print summary of op stats.
    36    void printSummary();
    37  
    38  private:
    39    llvm::StringMap<int64_t> opCount;
    40    llvm::raw_ostream &os;
    41  };
    42  } // namespace
    43  
    44  void PrintOpStatsPass::runOnModule() {
    45    opCount.clear();
    46  
    47    // Compute the operation statistics for each function in the module.
    48    for (auto &op : getModule())
    49      op.walk([&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
    50    printSummary();
    51  }
    52  
    53  void PrintOpStatsPass::printSummary() {
    54    os << "Operations encountered:\n";
    55    os << "-----------------------\n";
    56    SmallVector<StringRef, 64> sorted(opCount.keys());
    57    llvm::sort(sorted);
    58  
    59    // Split an operation name from its dialect prefix.
    60    auto splitOperationName = [](StringRef opName) {
    61      auto splitName = opName.split('.');
    62      return splitName.second.empty() ? std::make_pair("", splitName.first)
    63                                      : splitName;
    64    };
    65  
    66    // Compute the largest dialect and operation name.
    67    StringRef dialectName, opName;
    68    size_t maxLenOpName = 0, maxLenDialect = 0;
    69    for (const auto &key : sorted) {
    70      std::tie(dialectName, opName) = splitOperationName(key);
    71      maxLenDialect = std::max(maxLenDialect, dialectName.size());
    72      maxLenOpName = std::max(maxLenOpName, opName.size());
    73    }
    74  
    75    for (const auto &key : sorted) {
    76      std::tie(dialectName, opName) = splitOperationName(key);
    77  
    78      // Left-align the names (aligning on the dialect) and right-align the count
    79      // below. The alignment is for readability and does not affect CSV/FileCheck
    80      // parsing.
    81      if (dialectName.empty())
    82        os.indent(maxLenDialect + 3);
    83      else
    84        os << llvm::right_justify(dialectName, maxLenDialect + 2) << '.';
    85  
    86      // Left justify the operation name.
    87      os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
    88         << '\n';
    89    }
    90  }
    91  
    92  static PassRegistration<PrintOpStatsPass>
    93      pass("print-op-stats", "Print statistics of operations");