Commit fa9dd833 authored by River Riddle's avatar River Riddle
Browse files

[mlir] Refactor ModuleState into AsmState and expose it to users. 

Summary:
This allows for users to cache printer state, which can be costly to recompute. Each of the IR print methods gain a new overload taking this new state class.

Depends On D72293

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D72294
parent 23058f9d

mlir/include/mlir/IR/AsmState.h

0 → 100644
+52 −0
Original line number Diff line number Diff line 
//===- AsmState.h - State class for AsmPrinter ------------------*- C++ -*-===//

//

// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.

// See https://llvm.org/LICENSE.txt for license information.

// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

//

//===----------------------------------------------------------------------===//

//

// This file defines the AsmState class.

//

//===----------------------------------------------------------------------===//



#ifndef MLIR_IR_ASMSTATE_H_

#define MLIR_IR_ASMSTATE_H_



#include <memory>



namespace mlir {

class Operation;



namespace detail {

class AsmStateImpl;

} // end namespace detail



/// This class provides management for the lifetime of the state used when

/// printing the IR. It allows for alleviating the cost of recomputing the

/// internal state of the asm printer.

///

/// The IR should not be mutated in-between invocations using this state, and

/// the IR being printed must not be an parent of the IR originally used to

/// initialize this state. This means that if a child operation is provided, a

/// parent operation cannot reuse this state.

class AsmState {

public:

  /// Initialize the asm state at the level of the given operation.

  AsmState(Operation *op);

  ~AsmState();



  /// Return an instance of the internal implementation. Returns nullptr if the

  /// state has not been initialized.

  detail::AsmStateImpl &getImpl() { return *impl; }



private:

  AsmState() = delete;



  /// A pointer to allocated storage for the impl state.

  std::unique_ptr<detail::AsmStateImpl> impl;

};



} // end namespace mlir



#endif // MLIR_IR_ASMSTATE_H_

mlir/include/mlir/IR/Block.h

+2 −0
Original line number Diff line number Diff line
@@ -312,12 +312,14 @@ public: 
  }



  void print(raw_ostream &os);

  void print(raw_ostream &os, AsmState &state);

  void dump();



  /// Print out the name of the block without printing its body.

  /// NOTE: The printType argument is ignored.  We keep it for compatibility

  /// with LLVM dominator machinery that expects it to exist.

  void printAsOperand(raw_ostream &os, bool printType = true);

  void printAsOperand(raw_ostream &os, AsmState &state);



private:

  /// Pair of the parent object that owns this block and a bit that signifies if

mlir/include/mlir/IR/Module.h

+2 −0
Original line number Diff line number Diff line
@@ -57,6 +57,8 @@ public: 


  /// Print the this module in the custom top-level form.

  void print(raw_ostream &os, OpPrintingFlags flags = llvm::None);

  void print(raw_ostream &os, AsmState &state,

             OpPrintingFlags flags = llvm::None);

  void dump();



  //===--------------------------------------------------------------------===//

mlir/include/mlir/IR/OpDefinition.h

+4 −0
Original line number Diff line number Diff line
@@ -122,6 +122,10 @@ public: 
  void print(raw_ostream &os, OpPrintingFlags flags = llvm::None) {

    state->print(os, flags);

  }

  void print(raw_ostream &os, AsmState &asmState,

             OpPrintingFlags flags = llvm::None) {

    state->print(os, asmState, flags);

  }



  /// Dump this operation.

  void dump() { state->dump(); }

mlir/include/mlir/IR/Operation.h

+2 −0
Original line number Diff line number Diff line
@@ -187,6 +187,8 @@ public: 
  bool isBeforeInBlock(Operation *other);



  void print(raw_ostream &os, OpPrintingFlags flags = llvm::None);

  void print(raw_ostream &os, AsmState &state,

             OpPrintingFlags flags = llvm::None);

  void dump();



  //===--------------------------------------------------------------------===//