[mlir] Allow passing AsmState when printing Attributes and Types

[rriddle]

Summary

This allows for extracting assembly information when printing an attribute
or type, such as the dialect resources referenced. This functionality is used in
a followup that adds resource support to the bytecode. This change also results
in a nice cleanup of AsmPrinter now that we don't need to awkwardly workaround
optional AsmStates.

[jpienaar]

Would this also allow be propagated down to custom printers that want to access AsmState/OpPrintingFlags?

[tpopp]

Drive by warning. If this patch is depending on a non-landed commit, this comment might be wrong. It looks like one python binding test case will fail with this patch:

Exit Code: 2

Command Output (stderr):
--
<mlir_parser_buffer>:1:1: error: expected non-function type
BAD_TYPE_DOES_NOT_EXIST
^
error: vector elements must be int/index/float type but got 'none'
error: invalid tensor element type: 'none'
error: invalid tensor element type: 'none'
...
********************
Failed Tests (1):
  MLIR :: python/ir/builtin_types.py

This is only tested with cmake ... -DMLIR_ENABLE_BINDINGS_PYTHON=On and has dependencies on python -m pip install -r mlir/python/requirements.txt

[rriddle]

In D132728#3760082, @jpienaar wrote:

Would this also allow be propagated down to custom printers that want to access AsmState/OpPrintingFlags?

Yeah, it should make it a bit easier to propagate stuff along.

In D132728#3760715, @tpopp wrote:

Drive by warning. If this patch is depending on a non-landed commit, this comment might be wrong. It looks like one python binding test case will fail with this patch:

Exit Code: 2

Command Output (stderr):
--
<mlir_parser_buffer>:1:1: error: expected non-function type
BAD_TYPE_DOES_NOT_EXIST
^
error: vector elements must be int/index/float type but got 'none'
error: invalid tensor element type: 'none'
error: invalid tensor element type: 'none'
...
********************
Failed Tests (1):
  MLIR :: python/ir/builtin_types.py

This is only tested with cmake ... -DMLIR_ENABLE_BINDINGS_PYTHON=On and has dependencies on python -m pip install -r mlir/python/requirements.txt

Thanks for this, fixed.

mlir/include/mlir/IR/AsmState.h

 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LLVM.h"
 
 #include <memory>
 
 namespace mlir {
 class AsmResourcePrinter;
+class AsmDialectResourceHandle;
 class Operation;
 
 namespace detail {
 class AsmStateImpl;
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
 // Resources
public:
   /// and column in the output stream.
   using LocationMap = DenseMap<Operation *, std::pair<unsigned, unsigned>>;
 
   /// Initialize the asm state at the level of the given operation. A location
   /// map may optionally be provided to be populated when printing.
   AsmState(Operation *op,
            const OpPrintingFlags &printerFlags = OpPrintingFlags(),
            LocationMap *locationMap = nullptr);
+  AsmState(MLIRContext *ctx,
+           const OpPrintingFlags &printerFlags = OpPrintingFlags(),
+           LocationMap *locationMap = nullptr);
   ~AsmState();
 
   /// Get the printer flags.
   const OpPrintingFlags &getPrinterFlags() const;
 
   /// Return an instance of the internal implementation. Returns nullptr if the
   /// state has not been initialized.
   detail::AsmStateImpl &getImpl() { return *impl; }
public:
   template <typename CallableT>
   std::enable_if_t<std::is_convertible<
       CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value>
   attachResourcePrinter(StringRef name, CallableT &&printFn) {
     attachResourcePrinter(AsmResourcePrinter::fromCallable(
         name, std::forward<CallableT>(printFn)));
   }
 
+  /// Returns a map of dialect resources that were referenced when using this
+  /// state to print IR.
+  DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> &
+  getDialectResources() const;
+
 private:
   AsmState() = delete;
 
   /// A pointer to allocated storage for the impl state.
   std::unique_ptr<detail::AsmStateImpl> impl;
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/Attributes.h

 //===- Attributes.h - MLIR Attribute Classes --------------------*- C++ -*-===//
 //
 // Part of the LLVM 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ATTRIBUTES_H
 #define MLIR_IR_ATTRIBUTES_H
 
 #include "mlir/IR/AttributeSupport.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
+class AsmState;
 class StringAttr;
 
 /// Attributes are known-constant values of operations.
 ///
 /// Instances of the Attribute class are references to immortal key-value pairs
 /// with immutable, uniqued keys owned by MLIRContext. As such, an Attribute is
 /// a thin wrapper around an underlying storage pointer. Attributes are usually
 /// passed by value.
public:
 
   /// Get the dialect this attribute is registered to.
   Dialect &getDialect() const {
     return impl->getAbstractAttribute().getDialect();
   }
 
   /// Print the attribute.
   void print(raw_ostream &os) const;
+  void print(raw_ostream &os, AsmState &state) const;
   void dump() const;
 
   /// Get an opaque pointer to the attribute.
   const void *getAsOpaquePointer() const { return impl; }
   /// Construct an attribute from the opaque pointer representation.
   static Attribute getFromOpaquePointer(const void *ptr) {
     return Attribute(reinterpret_cast<const ImplType *>(ptr));
   }

mlir/include/mlir/IR/Types.h

 #define MLIR_IR_TYPES_H
 
 #include "mlir/IR/TypeSupport.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
+class AsmState;
+
 /// Instances of the Type class are uniqued, have an immutable identifier and an
 /// optional mutable component.  They wrap a pointer to the storage object owned
 /// by MLIRContext.  Therefore, instances of Type are passed around by value.
 ///
 /// Some types are "primitives" meaning they do not have any parameters, for
 /// example the Index type.  Parametric types have additional information that
 /// differentiates the types of the same class, for example the Integer type has
 /// bitwidth, making i8 and i16 belong to the same kind by be different
public:
   /// Return true if this is an integer (of any signedness) or a float type.
   bool isIntOrFloat() const;
   /// Return true if this is an integer (of any signedness), index, or float
   /// type.
   bool isIntOrIndexOrFloat() const;
 
   /// Print the current type.
   void print(raw_ostream &os) const;
+  void print(raw_ostream &os, AsmState &state) const;
   void dump() const;
 
   friend ::llvm::hash_code hash_value(Type arg);
 
   /// Methods for supporting PointerLikeTypeTraits.
   const void *getAsOpaquePointer() const {
     return static_cast<const void *>(impl);
   }

mlir/lib/IR/AsmPrinter.cpp

 
 /// This class manages the state of SSA value names.
 class SSANameState {
 public:
   /// A sentinel value used for values with names set.
   enum : unsigned { NameSentinel = ~0U };
 
   SSANameState(Operation *op, const OpPrintingFlags &printerFlags);
+  SSANameState() = default;
 
   /// Print the SSA identifier for the given value to 'stream'. If
   /// 'printResultNo' is true, it also presents the result number ('#' number)
   /// of this value.
   void printValueID(Value value, bool printResultNo, raw_ostream &stream) const;
 
   /// Print the operation identifier.
   void printOperationID(Operation *op, raw_ostream &stream) const;

 namespace mlir {
 namespace detail {
 class AsmStateImpl {
 public:
   explicit AsmStateImpl(Operation *op, const OpPrintingFlags &printerFlags,
                         AsmState::LocationMap *locationMap)
       : interfaces(op->getContext()), nameState(op, printerFlags),
         printerFlags(printerFlags), locationMap(locationMap) {}
+  explicit AsmStateImpl(MLIRContext *ctx, const OpPrintingFlags &printerFlags,
+                        AsmState::LocationMap *locationMap)
+      : interfaces(ctx), printerFlags(printerFlags), locationMap(locationMap) {}
 
   /// Initialize the alias state to enable the printing of aliases.
   void initializeAliases(Operation *op) {
     aliasState.initialize(op, printerFlags, interfaces);
   }
 
   /// Get the state used for aliases.
   AliasState &getAliasState() { return aliasState; }
public:
 
   /// Register the location, line and column, within the buffer that the given
   /// operation was printed at.
   void registerOperationLocation(Operation *op, unsigned line, unsigned col) {
     if (locationMap)
       (*locationMap)[op] = std::make_pair(line, col);
   }
 
+  /// Return the referenced dialect resources within the printer.
+  DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> &
+  getDialectResources() {
+    return dialectResources;
+  }
+
 private:
   /// Collection of OpAsm interfaces implemented in the context.
   DialectInterfaceCollection<OpAsmDialectInterface> interfaces;
 
   /// A collection of non-dialect resource printers.
   SmallVector<std::unique_ptr<AsmResourcePrinter>> externalResourcePrinters;
 
+  /// A set of dialect resources that were referenced during printing.
+  DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> dialectResources;
+
   /// The state used for attribute and type aliases.
   AliasState aliasState;
 
   /// The state used for SSA value names.
   SSANameState nameState;
 
   /// Flags that control op output.
   OpPrintingFlags printerFlags;
static OpPrintingFlags verifyOpAndAdjustFlags(Operation *op,
 
   return printerFlags;
 }
 
 AsmState::AsmState(Operation *op, const OpPrintingFlags &printerFlags,
                    LocationMap *locationMap)
     : impl(std::make_unique<AsmStateImpl>(
           op, verifyOpAndAdjustFlags(op, printerFlags), locationMap)) {}
+AsmState::AsmState(MLIRContext *ctx, const OpPrintingFlags &printerFlags,
+                   LocationMap *locationMap)
+    : impl(std::make_unique<AsmStateImpl>(ctx, printerFlags, locationMap)) {}
 AsmState::~AsmState() = default;
 
 const OpPrintingFlags &AsmState::getPrinterFlags() const {
   return impl->getPrinterFlags();
 }
 
 void AsmState::attachResourcePrinter(
     std::unique_ptr<AsmResourcePrinter> printer) {
   impl->externalResourcePrinters.emplace_back(std::move(printer));
 }
 
+DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> &
+AsmState::getDialectResources() const {
+  return impl->getDialectResources();
+}
+
 //===----------------------------------------------------------------------===//
 // AsmPrinter::Impl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 class AsmPrinter::Impl {
 public:
-  Impl(raw_ostream &os, OpPrintingFlags flags = llvm::None,
-       AsmStateImpl *state = nullptr)
-      : os(os), printerFlags(flags), state(state) {}
-  explicit Impl(Impl &other)
-      : Impl(other.os, other.printerFlags, other.state) {}
+  Impl(raw_ostream &os, AsmStateImpl &state)
+      : os(os), state(state), printerFlags(state.getPrinterFlags()) {}
+  explicit Impl(Impl &other) : Impl(other.os, other.state) {}
 
   /// Returns the output stream of the printer.
   raw_ostream &getStream() { return os; }
 
   template <typename Container, typename UnaryFunctor>
   inline void interleaveComma(const Container &c, UnaryFunctor eachFn) const {
     llvm::interleaveComma(c, os, eachFn);
   }
public:
   /// allows for the internal location to use an attribute alias.
   void printLocation(LocationAttr loc, bool allowAlias = false);
 
   /// Print a reference to the given resource that is owned by the given
   /// dialect.
   void printResourceHandle(const AsmDialectResourceHandle &resource) {
     auto *interface = cast<OpAsmDialectInterface>(resource.getDialect());
     os << interface->getResourceKey(resource);
-    dialectResources[resource.getDialect()].insert(resource);
+    state.getDialectResources()[resource.getDialect()].insert(resource);
   }
 
   void printAffineMap(AffineMap map);
   void
   printAffineExpr(AffineExpr expr,
                   function_ref<void(unsigned, bool)> printValueName = nullptr);
   void printAffineConstraint(AffineExpr expr, bool isEq);
   void printIntegerSet(IntegerSet set);
protected:
   };
   void printAffineExprInternal(
       AffineExpr expr, BindingStrength enclosingTightness,
       function_ref<void(unsigned, bool)> printValueName = nullptr);
 
   /// The output stream for the printer.
   raw_ostream &os;
 
+  /// An underlying assembly printer state.
+  AsmStateImpl &state;
+
   /// A set of flags to control the printer's behavior.
   OpPrintingFlags printerFlags;
 
-  /// An optional printer state for the module.
-  AsmStateImpl *state;
-
   /// A tracker for the number of new lines emitted during printing.
   NewLineCounter newLine;
-
-  /// A set of dialect resources that were referenced during printing.
-  DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> dialectResources;
 };
 } // namespace mlir
 
 void AsmPrinter::Impl::printTrailingLocation(Location loc, bool allowAlias) {
   // Check to see if we are printing debug information.
   if (!printerFlags.shouldPrintDebugInfo())
     return;
 
static void printFloatValue(const APFloat &apValue, raw_ostream &os) {
   os << str;
 }
 
 void AsmPrinter::Impl::printLocation(LocationAttr loc, bool allowAlias) {
   if (printerFlags.shouldPrintDebugInfoPrettyForm())
     return printLocationInternal(loc, /*pretty=*/true);
 
   os << "loc(";
-  if (!allowAlias || !state || failed(state->getAliasState().getAlias(loc, os)))
+  if (!allowAlias || failed(printAlias(loc)))
     printLocationInternal(loc);
   os << ')';
 }
 
 /// Returns true if the given dialect symbol data is simple enough to print in
 /// the pretty form. This is essentially when the symbol takes the form:
 ///   identifier (`<` body `>`)?
 static bool isDialectSymbolSimpleEnoughForPrettyForm(StringRef symName) {

 //
 // We choose to use a dense resource ElementsAttr literal with conspicuous
 // content to hopefully alert readers to the fact that this has been elided.
 static void printElidedElementsAttr(raw_ostream &os) {
   os << R"(dense_resource<__elided__>)";
 }
 
 LogicalResult AsmPrinter::Impl::printAlias(Attribute attr) {
-  return success(state && succeeded(state->getAliasState().getAlias(attr, os)));
+  return state.getAliasState().getAlias(attr, os);
 }
 
 LogicalResult AsmPrinter::Impl::printAlias(Type type) {
-  return success(state && succeeded(state->getAliasState().getAlias(type, os)));
+  return state.getAliasState().getAlias(type, os);
 }
 
 void AsmPrinter::Impl::printAttribute(Attribute attr,
                                       AttrTypeElision typeElision) {
   if (!attr) {
     os << "<<NULL ATTRIBUTE>>";
     return;
   }

 
 void AsmPrinter::Impl::printType(Type type) {
   if (!type) {
     os << "<<NULL TYPE>>";
     return;
   }
 
   // Try to print an alias for this type.
-  if (state && succeeded(state->getAliasState().getAlias(type, os)))
+  if (succeeded(printAlias(type)))
     return;
 
   TypeSwitch<Type>(type)
       .Case<OpaqueType>([&](OpaqueType opaqueTy) {
         printDialectSymbol(os, "!", opaqueTy.getDialectNamespace(),
                            opaqueTy.getTypeData());
       })
       .Case<IndexType>([&](Type) { os << "index"; })

 
 void AsmPrinter::Impl::printDialectAttribute(Attribute attr) {
   auto &dialect = attr.getDialect();
 
   // Ask the dialect to serialize the attribute to a string.
   std::string attrName;
   {
     llvm::raw_string_ostream attrNameStr(attrName);
-    Impl subPrinter(attrNameStr, printerFlags, state);
+    Impl subPrinter(attrNameStr, state);
     DialectAsmPrinter printer(subPrinter);
     dialect.printAttribute(attr, printer);
-
-    // FIXME: Delete this when we no longer require a nested printer.
-    for (auto &it : subPrinter.dialectResources)
-      for (const auto &resource : it.second)
-        dialectResources[it.first].insert(resource);
   }
   printDialectSymbol(os, "#", dialect.getNamespace(), attrName);
 }
 
 void AsmPrinter::Impl::printDialectType(Type type) {
   auto &dialect = type.getDialect();
 
   // Ask the dialect to serialize the type to a string.
   std::string typeName;
   {
     llvm::raw_string_ostream typeNameStr(typeName);
-    Impl subPrinter(typeNameStr, printerFlags, state);
+    Impl subPrinter(typeNameStr, state);
     DialectAsmPrinter printer(subPrinter);
     dialect.printType(type, printer);
-
-    // FIXME: Delete this when we no longer require a nested printer.
-    for (auto &it : subPrinter.dialectResources)
-      for (const auto &resource : it.second)
-        dialectResources[it.first].insert(resource);
   }
   printDialectSymbol(os, "!", dialect.getNamespace(), typeName);
 }
 
 void AsmPrinter::Impl::printEscapedString(StringRef str) {
   os << "\"";
   llvm::printEscapedString(str, os);
   os << "\"";

 namespace {
 /// This class contains the logic for printing operations, regions, and blocks.
 class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter {
 public:
   using Impl = AsmPrinter::Impl;
   using Impl::printType;
 
   explicit OperationPrinter(raw_ostream &os, AsmStateImpl &state)
-      : Impl(os, state.getPrinterFlags(), &state),
-        OpAsmPrinter(static_cast<Impl &>(*this)) {}
+      : Impl(os, state), OpAsmPrinter(static_cast<Impl &>(*this)) {}
 
   /// Print the given top-level operation.
   void printTopLevelOperation(Operation *op);
 
   /// Print the given operation with its indent and location.
   void print(Operation *op);
   /// Print the bare location, not including indentation/location/etc.
   void printOperation(Operation *op);
public:
   void printRegion(Region &region, bool printEntryBlockArgs,
                    bool printBlockTerminators, bool printEmptyBlock) override;
 
   /// Renumber the arguments for the specified region to the same names as the
   /// SSA values in namesToUse. This may only be used for IsolatedFromAbove
   /// operations. If any entry in namesToUse is null, the corresponding
   /// argument name is left alone.
   void shadowRegionArgs(Region &region, ValueRange namesToUse) override {
-    state->getSSANameState().shadowRegionArgs(region, namesToUse);
+    state.getSSANameState().shadowRegionArgs(region, namesToUse);
   }
 
   /// Print the given affine map with the symbol and dimension operands printed
   /// inline with the map.
   void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
                               ValueRange operands) override;
 
   /// Print the given affine expression with the symbol and dimension operands
private:
 
   // This is the current indentation level for nested structures.
   unsigned currentIndent = 0;
 };
 } // namespace
 
 void OperationPrinter::printTopLevelOperation(Operation *op) {
   // Output the aliases at the top level that can't be deferred.
-  state->getAliasState().printNonDeferredAliases(os, newLine);
+  state.getAliasState().printNonDeferredAliases(os, newLine);
 
   // Print the module.
   print(op);
   os << newLine;
 
   // Output the aliases at the top level that can be deferred.
-  state->getAliasState().printDeferredAliases(os, newLine);
+  state.getAliasState().printDeferredAliases(os, newLine);
 
   // Output any file level metadata.
   printFileMetadataDictionary(op);
 }
 
 void OperationPrinter::printFileMetadataDictionary(Operation *op) {
   bool sawMetadataEntry = false;
   auto checkAddMetadataDict = [&] {
auto processProvider = [&](StringRef dictName, StringRef name, auto &provider,
     provider.buildResources(op, providerArgs..., entryBuilder);
 
     if (hadEntry)
       os << newLine << "    }";
   };
 
   // Print the `dialect_resources` section if we have any dialects with
   // resources.
-  for (const OpAsmDialectInterface &interface : state->getDialectInterfaces()) {
+  for (const OpAsmDialectInterface &interface : state.getDialectInterfaces()) {
+    auto &dialectResources = state.getDialectResources();
     StringRef name = interface.getDialect()->getNamespace();
     auto it = dialectResources.find(interface.getDialect());
     if (it != dialectResources.end())
       processProvider("dialect", name, interface, it->second);
     else
       processProvider("dialect", name, interface,
                       SetVector<AsmDialectResourceHandle>());
   }
   if (hadResource)
     os << newLine << "  }";
 
   // Print the `external_resources` section if we have any external clients with
   // resources.
   hadResource = false;
-  for (const auto &printer : state->getResourcePrinters())
+  for (const auto &printer : state.getResourcePrinters())
     processProvider("external", printer.getName(), printer);
   if (hadResource)
     os << newLine << "  }";
 }
 
 /// Print a block argument in the usual format of:
 ///   %ssaName : type {attr1=42} loc("here")
 /// where location printing is controlled by the standard internal option.
void OperationPrinter::printRegionArgument(BlockArgument arg,
   }
   printOptionalAttrDict(argAttrs);
   // TODO: We should allow location aliases on block arguments.
   printTrailingLocation(arg.getLoc(), /*allowAlias*/ false);
 }
 
 void OperationPrinter::print(Operation *op) {
   // Track the location of this operation.
-  state->registerOperationLocation(op, newLine.curLine, currentIndent);
+  state.registerOperationLocation(op, newLine.curLine, currentIndent);
 
   os.indent(currentIndent);
   printOperation(op);
   printTrailingLocation(op->getLoc());
   if (printerFlags.shouldPrintValueUsers())
     printUsersComment(op);
 }
 
 void OperationPrinter::printOperation(Operation *op) {
   if (size_t numResults = op->getNumResults()) {
     auto printResultGroup = [&](size_t resultNo, size_t resultCount) {
       printValueID(op->getResult(resultNo), /*printResultNo=*/false);
       if (resultCount > 1)
         os << ':' << resultCount;
     };
 
     // Check to see if this operation has multiple result groups.
-    ArrayRef<int> resultGroups = state->getSSANameState().getOpResultGroups(op);
+    ArrayRef<int> resultGroups = state.getSSANameState().getOpResultGroups(op);
     if (!resultGroups.empty()) {
       // Interleave the groups excluding the last one, this one will be handled
       // separately.
       interleaveComma(llvm::seq<int>(0, resultGroups.size() - 1), [&](int i) {
         printResultGroup(resultGroups[i],
                          resultGroups[i + 1] - resultGroups[i]);
       });
       os << ", ";
void OperationPrinter::printGenericOp(Operation *op, bool printOpName) {
   printOptionalAttrDict(attrs);
 
   // Print the type signature of the operation.
   os << " : ";
   printFunctionalType(op);
 }
 
 void OperationPrinter::printBlockName(Block *block) {
-  os << state->getSSANameState().getBlockInfo(block).name;
+  os << state.getSSANameState().getBlockInfo(block).name;
 }
 
 void OperationPrinter::print(Block *block, bool printBlockArgs,
                              bool printBlockTerminator) {
   // Print the block label and argument list if requested.
   if (printBlockArgs) {
     os.indent(currentIndent);
     printBlockName(block);
if (printBlockArgs) {
     } else if (auto *pred = block->getSinglePredecessor()) {
       os << "  // pred: ";
       printBlockName(pred);
     } else {
       // We want to print the predecessors in a stable order, not in
       // whatever order the use-list is in, so gather and sort them.
       SmallVector<BlockInfo, 4> predIDs;
       for (auto *pred : block->getPredecessors())
-        predIDs.push_back(state->getSSANameState().getBlockInfo(pred));
+        predIDs.push_back(state.getSSANameState().getBlockInfo(pred));
       llvm::sort(predIDs, [](BlockInfo lhs, BlockInfo rhs) {
         return lhs.ordering < rhs.ordering;
       });
 
       os << "  // " << predIDs.size() << " preds: ";
 
       interleaveComma(predIDs, [&](BlockInfo pred) { os << pred.name; });
     }
for (auto &op : range) {
     print(&op);
     os << newLine;
   }
   currentIndent -= indentWidth;
 }
 
 void OperationPrinter::printValueID(Value value, bool printResultNo,
                                     raw_ostream *streamOverride) const {
-  state->getSSANameState().printValueID(value, printResultNo,
-                                        streamOverride ? *streamOverride : os);
+  state.getSSANameState().printValueID(value, printResultNo,
 }
 
 void OperationPrinter::printOperationID(Operation *op,
                                         raw_ostream *streamOverride) const {
-  state->getSSANameState().printOperationID(op, streamOverride ? *streamOverride
-                                                               : os);
+  state.getSSANameState().printOperationID(op, streamOverride ? *streamOverride
 }
 
 void OperationPrinter::printSuccessor(Block *successor) {
   printBlockName(successor);
 }
 
 void OperationPrinter::printSuccessorAndUseList(Block *successor,
                                                 ValueRange succOperands) {
void OperationPrinter::printAffineExprOfSSAIds(AffineExpr expr,
   printAffineExpr(expr, printValueName);
 }
 
 //===----------------------------------------------------------------------===//
 // print and dump methods
 //===----------------------------------------------------------------------===//
 
 void Attribute::print(raw_ostream &os) const {
-  AsmPrinter::Impl(os).printAttribute(*this);
+  if (!*this) {
+    os << "<<NULL ATTRIBUTE>>";
+    return;
+  }
+
+  AsmState state(getContext());
+  print(os, state);
+}
+void Attribute::print(raw_ostream &os, AsmState &state) const {
+  AsmPrinter::Impl(os, state.getImpl()).printAttribute(*this);
 }
 
 void Attribute::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
 void Type::print(raw_ostream &os) const {
-  AsmPrinter::Impl(os).printType(*this);
+  if (!*this) {
+    os << "<<NULL TYPE>>";
+    return;
+  }
+
+  AsmState state(getContext());
+  print(os, state);
+}
+void Type::print(raw_ostream &os, AsmState &state) const {
+  AsmPrinter::Impl(os, state.getImpl()).printType(*this);
 }
 
 void Type::dump() const { print(llvm::errs()); }
 
 void AffineMap::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
 void IntegerSet::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
 void AffineExpr::print(raw_ostream &os) const {
   if (!expr) {
     os << "<<NULL AFFINE EXPR>>";
     return;
   }
-  AsmPrinter::Impl(os).printAffineExpr(*this);
+  AsmState state(getContext());
+  AsmPrinter::Impl(os, state.getImpl()).printAffineExpr(*this);
 }
 
 void AffineExpr::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
 void AffineMap::print(raw_ostream &os) const {
   if (!map) {
     os << "<<NULL AFFINE MAP>>";
     return;
   }
-  AsmPrinter::Impl(os).printAffineMap(*this);
+  AsmState state(getContext());
+  AsmPrinter::Impl(os, state.getImpl()).printAffineMap(*this);
 }
 
 void IntegerSet::print(raw_ostream &os) const {
-  AsmPrinter::Impl(os).printIntegerSet(*this);
+  AsmState state(getContext());
+  AsmPrinter::Impl(os, state.getImpl()).printIntegerSet(*this);
 }
 
 void Value::print(raw_ostream &os) { print(os, OpPrintingFlags()); }
 void Value::print(raw_ostream &os, const OpPrintingFlags &flags) {
   if (!impl) {
     os << "<<NULL VALUE>>";
     return;
   }