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

     1  //===- RemoveInstrumentationPass.cpp - Removes instrumentation ------------===//
     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  // This file defines a pass to remove any instrumentation ops. It is often one
    19  // of the final steps when performing quantization and is run after any
    20  // decisions requiring instrumentation have been made.
    21  //
    22  //===----------------------------------------------------------------------===//
    23  
    24  #include "mlir/Dialect/QuantOps/QuantOps.h"
    25  #include "mlir/IR/MLIRContext.h"
    26  #include "mlir/IR/PatternMatch.h"
    27  #include "mlir/Pass/Pass.h"
    28  #include "mlir/Quantizer/Transforms/Passes.h"
    29  
    30  using namespace mlir;
    31  using namespace mlir::quantizer;
    32  using namespace mlir::quant;
    33  
    34  namespace {
    35  
    36  class RemoveInstrumentationPass
    37      : public FunctionPass<RemoveInstrumentationPass> {
    38    void runOnFunction() override;
    39  };
    40  
    41  template <typename OpTy>
    42  class RemoveIdentityOpRewrite : public RewritePattern {
    43  public:
    44    RemoveIdentityOpRewrite(MLIRContext *context)
    45        : RewritePattern(OpTy::getOperationName(), 1, context) {}
    46  
    47    PatternMatchResult matchAndRewrite(Operation *op,
    48                                       PatternRewriter &rewriter) const override {
    49      assert(op->getNumOperands() == 1);
    50      assert(op->getNumResults() == 1);
    51  
    52      rewriter.replaceOp(op, op->getOperand(0));
    53      return matchSuccess();
    54    }
    55  };
    56  
    57  } // end anonymous namespace
    58  
    59  void RemoveInstrumentationPass::runOnFunction() {
    60    OwningRewritePatternList patterns;
    61    auto func = getFunction();
    62    auto *context = &getContext();
    63    patterns.insert<RemoveIdentityOpRewrite<StatisticsOp>,
    64                    RemoveIdentityOpRewrite<StatisticsRefOp>,
    65                    RemoveIdentityOpRewrite<CoupledRefOp>>(context);
    66    applyPatternsGreedily(func, patterns);
    67  }
    68  
    69  std::unique_ptr<FunctionPassBase>
    70  mlir::quantizer::createRemoveInstrumentationPass() {
    71    return std::make_unique<RemoveInstrumentationPass>();
    72  }
    73  
    74  static PassRegistration<RemoveInstrumentationPass>
    75      pass("quantizer-remove-instrumentation",
    76           "Removes instrumentation and hints which have no effect on final "
    77           "execution");