[mlir][bufferize][NFC] Rename DialectAnalysisState and move to OneShotAnalysis

[springerm]

Summary

DialectAnalysisState is now OneShotAnalysisState::Extension.

This state extension mechanism is needed only for One-Shot Analysis, so it is moved from BufferizableOpInterface.h to OneShotAnalysis.h.

Extensions are now identified via TypeIDs instead of StringRefs. The API of state extensions is cleaned up and follows the same pattern as other extension mechanisms in MLIR (e.g., transform::TransformState::Extension).

Also delete some dead code.

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[ftynse]

LGTM with comments addressed.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

408–409

It is dangerous to have a public constructor that takes TypeID because the caller can pass an arbitrary ID. Instead, consider adding an additional protected constructor and delegate to it. There's an even safer version of having an AnalysisBase CRTP class that handles all TypeID-related logic, but it may be too verbose for this use case.

415–417

Nit: there are some changes in how isa/dyn_cast are supposed to be set up (see https://llvm.org/docs/HowToSetUpLLVMStyleRTTI.html#advanced-use-cases), but I haven't followed them deeply enough to say if you want to use that instead of classof here.

426–427

You probably want to add MLIR_DECLARE/DEFINE_EXPLICIT_TYPE_ID macros for all relevant classes. See the documentation in TypeID.h

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

233

Nit: MLIR style has a newline after the template header.

244–263

One version is usually implemented by calling another with sufficient const_casting

[springerm]

rebase and address comments

[springerm]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

415–417

classof seems easier for such small (and fixed) hierarchies. (following the same implementation approach as BlockArgumentImpl::classof)

[springerm]

typo

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 
 namespace mlir {
 class OpBuilder;
 
 namespace bufferization {
 
 class AnalysisState;
 class BufferizableOpInterface;
-struct DialectAnalysisState;
 
 class OpFilter {
 public:
   /// An op filter entry. Filters can be used to specify which ops should be
   /// processed by the bufferization.
   struct Entry {
     /// If the filter function evaluates to `true`, the filter matches.
     using FilterFn = std::function<bool(Operation *)>;
struct BufferizationOptions {
   /// AllocatorFn.
   using DeallocationFn =
       std::function<LogicalResult(OpBuilder &, Location, Value)>;
   /// Memcpy function: Generate a memcpy between two buffers.
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
   /// Initializer function for analysis state.
   using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
-  /// Initializer function for dialect-specific analysis state.
-  using DialectStateInitFn =
-      std::function<std::unique_ptr<DialectAnalysisState>()>;
   /// Tensor -> MemRef type converter.
   /// Parameters: Value, memory space, bufferization options
   using UnknownTypeConverterFn = std::function<BaseMemRefType(
       Value, Attribute memorySpace, const BufferizationOptions &)>;
 
   BufferizationOptions();
 
   /// Try to cast the given op to BufferizableOpInterface if the op is allow
struct BufferizationOptions {
   bool printConflicts = false;
 
   /// Buffer alignment for new memory allocations.
   unsigned int bufferAlignment = 128;
 
   /// Initializer functions for analysis state. These can be used to
   /// initialize dialect-specific analysis state.
   SmallVector<AnalysisStateInitFn> stateInitializers;
-
-  /// Add a analysis state initializer that initializes the specified
-  /// dialect-specific analysis state.
-  void addDialectStateInitializer(StringRef name, const DialectStateInitFn &fn);
 };
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
 
 /// Return `true` if the given value is a BlockArgument of a func::FuncOp.
 bool isFunctionArgument(Value value);
 
-/// Dialect-specific analysis state. Analysis/bufferization information
-/// that is specific to ops from a certain dialect can be stored in derived
-/// variants of this struct.
-struct DialectAnalysisState {
-  DialectAnalysisState() = default;
-
-  virtual ~DialectAnalysisState() = default;
-
-  // Copying state is forbidden. Always pass as reference.
-  DialectAnalysisState(const DialectAnalysisState &) = delete;
-};
-
 /// AnalysisState provides a variety of helper functions for dealing with
 /// tensor values.
 class AnalysisState {
 public:
   /// Determine which OpOperand* will alias with `result` if the op is
   /// bufferized in place. Return an empty vector if the op is not bufferizable.
   SmallVector<OpOperand *> getAliasingOpOperand(OpResult result) const;
 
public:
 
   /// Return true if the given tensor (or an aliasing tensor) is yielded from
   /// the containing block. Also include all aliasing tensors in the same block.
   ///
   /// Note: In the absence of an analysis, an implementation may return true for
   /// any given tensor.
   virtual bool isTensorYielded(Value tensor) const;
 
-  /// Return `true` if the given dialect state exists.
-  bool hasDialectState(StringRef name) const {
-    auto it = dialectState.find(name);
-    return it != dialectState.end();
-  }
-
-  /// Return dialect-specific bufferization state.
-  template <typename StateT>
-  Optional<const StateT *> getDialectState(StringRef name) const {
-    auto it = dialectState.find(name);
-    if (it == dialectState.end())
-      return None;
-    return static_cast<const StateT *>(it->getSecond().get());
-  }
-
-  /// Return dialect-specific analysis state or create one if none exists.
-  template <typename StateT>
-  StateT &getOrCreateDialectState(StringRef name) {
-    // Create state if it does not exist yet.
-    if (!hasDialectState(name))
-      dialectState[name] = std::make_unique<StateT>();
-    return static_cast<StateT &>(*dialectState[name]);
-  }
-
-  void insertDialectState(StringRef name,
-                          std::unique_ptr<DialectAnalysisState> state) {
-    assert(!dialectState.count(name) && "dialect state already initialized");
-    dialectState[name] = std::move(state);
-  }
-
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }
 
-  explicit AnalysisState(const BufferizationOptions &options);
+  AnalysisState(const BufferizationOptions &options);
 
   // AnalysisState should be passed as a reference.
   AnalysisState(const AnalysisState &) = delete;
 
   virtual ~AnalysisState() = default;
 
-private:
-  /// Dialect-specific analysis state.
-  DenseMap<StringRef, std::unique_ptr<DialectAnalysisState>> dialectState;
+  static bool classof(const AnalysisState *base) { return true; }
+
+  TypeID getType() const { return type; }
+
+protected:
+  AnalysisState(const BufferizationOptions &options, TypeID type);
 
+private:
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
+
+  /// The type of analysis.
+  TypeID type;
 };
 
 /// Create an AllocTensorOp for the given shaped value (memref or tensor).
 /// If `copy` is set, the shaped value is copied. Otherwise, a tensor with
 /// undefined contents is allocated.
 FailureOr<Value>
 allocateTensorForShapedValue(OpBuilder &b, Location loc, Value shapedValue,
                              bool escape, const BufferizationOptions &options,

 /// places.
 bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
                                unsigned index);
 } // namespace detail
 
 } // namespace bufferization
 } // namespace mlir
 
+MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::bufferization::AnalysisState);
+
 //===----------------------------------------------------------------------===//
 // Bufferization Interfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h

 //===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
 //
 // 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_BUFFERIZATION_TRANSFORMS_FUNCBUFFERIZABLEOPINTERFACEIMPL_H
 #define MLIR_BUFFERIZATION_TRANSFORMS_FUNCBUFFERIZABLEOPINTERFACEIMPL_H
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 
 namespace mlir {
 class DialectRegistry;
 
 namespace func {
 class FuncOp;
 } // namespace func
 
 namespace bufferization {
 namespace func_ext {
 /// The state of analysis of a FuncOp.
 enum class FuncOpAnalysisState { NotAnalyzed, InProgress, Analyzed };
 
 using func::FuncOp;
 
 /// Extra analysis state that is required for bufferization of function
 /// boundaries.
-struct FuncAnalysisState : public DialectAnalysisState {
+struct FuncAnalysisState : public OneShotAnalysisState::Extension {
+  FuncAnalysisState(OneShotAnalysisState &state)
+      : OneShotAnalysisState::Extension(state) {}
+
   // Note: Function arguments and/or function return values may disappear during
   // bufferization. Functions and their CallOps are analyzed and bufferized
   // separately. To ensure that a CallOp analysis/bufferization can access an
   // already bufferized function's analysis results, we store bbArg/return value
   // indices instead of BlockArguments/OpOperand pointers.
 
   /// A set of block argument indices.
   using BbArgIndexSet = DenseSet<int64_t>;

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

 public:
   OneShotAnalysisState(Operation *op,
                        const OneShotBufferizationOptions &options);
 
   OneShotAnalysisState(const OneShotAnalysisState &) = delete;
 
   ~OneShotAnalysisState() override = default;
 
+  static bool classof(const AnalysisState *base) {
+    return base->getType() == TypeID::get<OneShotAnalysisState>();
+  }
+
   /// Return a reference to the BufferizationAliasInfo.
   BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
public:
 
   /// Return true if the buffer of the given tensor value is written to. Must
   /// not be called for values inside not yet analyzed functions.
   bool isValueWritten(Value value) const;
 
   /// Return true if the buffer of the given tensor value is writable.
   bool isWritable(Value value) const;
 
+  /// Base class for OneShotAnalysisState extensions that allow
+  /// OneShotAnalysisState to contain user-specified information in the state
+  /// object. Clients are expected to derive this class, add the desired fields,
+  /// and make the derived class compatible with the MLIR TypeID mechanism.
+  ///
+  /// ```mlir
+  /// class MyExtension final : public OneShotAnalysisState::Extension {
+  /// public:
+  ///   MyExtension(OneShotAnalysisState &state, int myData)
+  ///       : Extension(state) {...}
+  /// private:
+  ///   int mySupplementaryData;
+  /// };
+  /// ```
+  ///
+  /// Instances of this and derived classes are not expected to be created by
+  /// the user, instead they are directly constructed within a
+  /// OneShotAnalysisState. A OneShotAnalysisState can only contain one
+  /// extension with the given TypeID. Extensions can be obtained from a
+  /// OneShotAnalysisState instance.
+  ///
+  /// ```mlir
+  /// state.addExtension<MyExtension>(/*myData=*/42);
+  /// MyExtension *ext = state.getExtension<MyExtension>();
+  /// ext->doSomething();
+  /// ```
+  class Extension {
+    // Allow OneShotAnalysisState to allocate Extensions.
+    friend class OneShotAnalysisState;
+
+  public:
+    /// Base virtual destructor.
+    // Out-of-line definition ensures symbols are emitted in a single object
+    // file.
+    virtual ~Extension();
+
+  protected:
+    /// Constructs an extension of the given state object.
+    Extension(OneShotAnalysisState &state) : state(state) {}
+
+    /// Provides read-only access to the parent OneShotAnalysisState object.
+    const OneShotAnalysisState &getAnalysisState() const { return state; }
+
+  private:
+    /// Back-reference to the state that is being extended.
+    OneShotAnalysisState &state;
+  };
+
+  /// Adds a new Extension of the type specified as template parameter,
+  /// constructing it with the arguments provided. The extension is owned by the
+  /// OneShotAnalysisState. It is expected that the state does not already have
+  /// an extension of the same type. Extension constructors are expected to take
+  /// a reference to OneShotAnalysisState as first argument, automatically
+  /// supplied by this call.
+  template <typename Ty, typename... Args>
+  Ty &addExtension(Args &&...args) {
+    static_assert(
+        std::is_base_of<Extension, Ty>::value,
+        "only a class derived from OneShotAnalysisState::Extension is allowed");
+    auto ptr = std::make_unique<Ty>(*this, std::forward<Args>(args)...);
+    auto result = extensions.try_emplace(TypeID::get<Ty>(), std::move(ptr));
+    assert(result.second && "extension already added");
+    return *static_cast<Ty *>(result.first->second.get());
+  }
+
+  /// Returns the extension of the specified type.
+  template <typename Ty>
+  Ty *getExtension() {
+    static_assert(
+        std::is_base_of<Extension, Ty>::value,
+        "only a class derived from OneShotAnalysisState::Extension is allowed");
+    auto iter = extensions.find(TypeID::get<Ty>());
+    if (iter == extensions.end())
+      return nullptr;
+    return static_cast<Ty *>(iter->second.get());
+  }
+
+  /// Returns the extension of the specified type.
+  template <typename Ty>
+  const Ty *getExtension() const {
+    return const_cast<OneShotAnalysisState *>(this)->getExtension<Ty>();
+  }
+
 private:
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runOneShotBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// A set of all tensors (and maybe aliasing tensors) that yielded from a
   /// block.
   DenseSet<Value> yieldedTensors;
 
   /// A set of uses of tensors that have undefined contents.
   DenseSet<OpOperand *> undefinedTensorUses;
+
+  /// Extensions attached to the state, identified by the TypeID of their type.
+  /// Only one extension of any given type is allowed.
+  DenseMap<TypeID, std::unique_ptr<Extension>> extensions;
 };
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state);
 
 /// Run One-Shot Bufferize on the given op: Analysis + Bufferization
 LogicalResult runOneShotBufferize(Operation *op,
                                   const OneShotBufferizationOptions &options);
 
 } // namespace bufferization
 } // namespace mlir
 
+MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::bufferization::OneShotAnalysisState);
+
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

 namespace mlir {
 namespace bufferization {
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.cpp.inc"
 
 } // namespace bufferization
 } // namespace mlir
 
+MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::bufferization::AnalysisState);
+
 #define DEBUG_TYPE "bufferizable-op-interface"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << (X))
 
 using namespace mlir;
 using namespace bufferization;
 
 Operation *bufferization::getOwnerOfValue(Value value) {

 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Value value) const {
   if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
     if (isOpAllowed(bufferizableOp.getOperation()))
       return bufferizableOp;
   return nullptr;
 }
 
-void BufferizationOptions::addDialectStateInitializer(
-    StringRef name, const DialectStateInitFn &fn) {
-  stateInitializers.push_back(
-      [=](AnalysisState &state) { state.insertDialectState(name, fn()); });
-}
-
 //===----------------------------------------------------------------------===//
 // Helper functions for BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
 static void setInsertionPointAfter(OpBuilder &b, Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     b.setInsertionPointToStart(bbArg.getOwner());
   } else {
return findValueInReverseUseDefChain(value, [&](Value value) {
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return true;
     return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
   });
 }
 
 AnalysisState::AnalysisState(const BufferizationOptions &options)
-    : options(options) {
+    : AnalysisState(options, TypeID::get<AnalysisState>()) {}
+
+AnalysisState::AnalysisState(const BufferizationOptions &options, TypeID type)
+    : options(options), type(type) {
   for (const BufferizationOptions::AnalysisStateInitFn &fn :
        options.stateInitializers)
     fn(*this);
 }
 
 bool AnalysisState::canOmitTensorCopy(OpOperand &opOperand) const {
   // Do not copy if the tensor has undefined contents.
   if (hasUndefinedContents(&opOperand))

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

if (!sym)
     return nullptr;
   return dyn_cast_or_null<FuncOp>(
       SymbolTable::lookupNearestSymbolFrom(callOp, sym));
 }
 
 /// Get FuncAnalysisState.
 static const FuncAnalysisState &
 getFuncAnalysisState(const AnalysisState &state) {
-  Optional<const FuncAnalysisState *> maybeState =
-      state.getDialectState<FuncAnalysisState>(
-          func::FuncDialect::getDialectNamespace());
-  assert(maybeState && "FuncAnalysisState does not exist");
-  return **maybeState;
+  assert(isa<OneShotAnalysisState>(state) && "expected OneShotAnalysisState");
+  auto *result = static_cast<const OneShotAnalysisState &>(state)
+                     .getExtension<FuncAnalysisState>();
+  assert(result && "FuncAnalysisState does not exist");
+  return *result;
 }
 
 /// Return the state (phase) of analysis of the FuncOp.
 static FuncOpAnalysisState getFuncOpAnalysisState(const AnalysisState &state,
                                                   FuncOp funcOp) {
-  Optional<const FuncAnalysisState *> maybeState =
-      state.getDialectState<FuncAnalysisState>(
-          func::FuncDialect::getDialectNamespace());
-  if (!maybeState.has_value())
+  if (!isa<OneShotAnalysisState>(state))
     return FuncOpAnalysisState::NotAnalyzed;
-  const auto &analyzedFuncOps = maybeState.value()->analyzedFuncOps;
+  auto *funcState = static_cast<const OneShotAnalysisState &>(state)
+                        .getExtension<FuncAnalysisState>();
+  if (!funcState)
+    return FuncOpAnalysisState::NotAnalyzed;
+  const auto &analyzedFuncOps = funcState->analyzedFuncOps;
   auto it = analyzedFuncOps.find(funcOp);
   if (it == analyzedFuncOps.end())
     return FuncOpAnalysisState::NotAnalyzed;
   return it->second;
 }
 
 /// Return the index of the bbArg in the given FuncOp that is equivalent to the
 /// specified return value (if any).

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
+MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::bufferization::OneShotAnalysisState);
+
 // Run mlir-opt with `-debug-only="one-shot-analysis"` for detailed debug
 // output.
 #define DEBUG_TYPE "one-shot-analysis"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }

 }
 
 //===----------------------------------------------------------------------===//
 // OneShotAnalysisState
 //===----------------------------------------------------------------------===//
 
 OneShotAnalysisState::OneShotAnalysisState(
     Operation *op, const OneShotBufferizationOptions &options)
-    : AnalysisState(options), aliasInfo(op) {
+    : AnalysisState(options, TypeID::get<OneShotAnalysisState>()),
+      aliasInfo(op) {
   // Set up alias sets for OpResults that must bufferize in-place. This should
   // be done before making any other bufferization decisions.
   op->walk([&](BufferizableOpInterface bufferizableOp) {
     if (!options.isOpAllowed(bufferizableOp))
       return WalkResult::skip();
     for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
       if (opOperand.get().getType().isa<TensorType>())
         if (bufferizableOp.mustBufferizeInPlace(opOperand, *this)) {
if (auto bbArg = value.dyn_cast<BlockArgument>())
     if (auto bufferizableOp =
             getOptions().dynCastBufferizableOp(bbArg.getOwner()->getParentOp()))
       return bufferizableOp.isWritable(bbArg, *this);
 
   // Not a bufferizable op: The conservative answer is "not writable".
   return false;
 }
 
+OneShotAnalysisState::Extension::~Extension() = default;
+
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
                                  const BufferizationAliasInfo &aliasInfo,
                                  const AnalysisState &state) {

 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        OneShotAnalysisState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const auto &options =
       static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
-  // Catch incorrect API usage.
-  assert((state.hasDialectState(func::FuncDialect::getDialectNamespace()) ||
-          !options.bufferizeFunctionBoundaries) &&
-         "must use ModuleBufferize to bufferize function boundaries");
-
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo, state);

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::bufferization::func_ext;
 
 /// A mapping of FuncOps to their callers.
 using FuncCallerMap = DenseMap<func::FuncOp, DenseSet<Operation *>>;
 
-/// Get FuncAnalysisState.
-static const FuncAnalysisState &
-getFuncAnalysisState(const AnalysisState &state) {
-  Optional<const FuncAnalysisState *> maybeState =
-      state.getDialectState<FuncAnalysisState>(
-          func::FuncDialect::getDialectNamespace());
-  assert(maybeState && "FuncAnalysisState does not exist");
-  return **maybeState;
-}
-
 /// Get or create FuncAnalysisState.
-static FuncAnalysisState &getFuncAnalysisState(AnalysisState &state) {
-  return state.getOrCreateDialectState<FuncAnalysisState>(
-      func::FuncDialect::getDialectNamespace());
-}
-
-/// Return the state (phase) of analysis of the FuncOp.
-/// Used for debug modes.
-LLVM_ATTRIBUTE_UNUSED
-static FuncOpAnalysisState getFuncOpAnalysisState(const AnalysisState &state,
-                                                  func::FuncOp funcOp) {
-  const FuncAnalysisState &funcState = getFuncAnalysisState(state);
-  auto it = funcState.analyzedFuncOps.find(funcOp);
-  if (it == funcState.analyzedFuncOps.end())
-    return FuncOpAnalysisState::NotAnalyzed;
-  return it->second;
+static FuncAnalysisState &
+getOrCreateFuncAnalysisState(OneShotAnalysisState &state) {
+  auto *result = state.getExtension<FuncAnalysisState>();
+  if (result)
+    return *result;
+  return state.addExtension<FuncAnalysisState>();
 }
 
 /// Return the unique ReturnOp that terminates `funcOp`.
 /// Return nullptr if there is no such unique ReturnOp.
 static func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp) {
   func::ReturnOp returnOp;
   for (Block &b : funcOp.getBody()) {
     if (auto candidateOp = dyn_cast<func::ReturnOp>(b.getTerminator())) {
static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
   equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
   OpBuilder b(op->getContext());
   op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
 }
 
 /// Store function BlockArguments that are equivalent to/aliasing a returned
 /// value in FuncAnalysisState.
-static LogicalResult aliasingFuncOpBBArgsAnalysis(FuncOp funcOp,
-                                                  OneShotAnalysisState &state) {
-  FuncAnalysisState &funcState = getFuncAnalysisState(state);
-
+static LogicalResult
+aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
+                             FuncAnalysisState &funcState) {
   // Support only single return-terminated block in the function.
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   for (OpOperand &returnVal : returnOp->getOpOperands())
     if (returnVal.get().getType().isa<RankedTensorType>())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<RankedTensorType>()) {
if (isRead && isWritten) {
     accessType = b.getStringAttr("none");
   }
   funcOp.setArgAttr(bbArg.getArgNumber(), "bufferization.access", accessType);
 }
 
 /// Determine which FuncOp bbArgs are read and which are written. When run on a
 /// function with unknown ops, we conservatively assume that such ops bufferize
 /// to a read + write.
-static LogicalResult funcOpBbArgReadWriteAnalysis(FuncOp funcOp,
-                                                  OneShotAnalysisState &state) {
-  FuncAnalysisState &funcState = getFuncAnalysisState(state);
-
+static LogicalResult
+funcOpBbArgReadWriteAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
+                             FuncAnalysisState &funcState) {
   // If the function has no body, conservatively assume that all args are
   // read + written.
   if (funcOp.getBody().empty()) {
     for (BlockArgument bbArg : funcOp.getArguments()) {
       funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
       funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
     }
 

 }
 
 /// Gather equivalence info of CallOps.
 /// Note: This only adds new equivalence info if the called function was already
 /// analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(func::FuncOp funcOp,
                                 BufferizationAliasInfo &aliasInfo,
-                                OneShotAnalysisState &state) {
-  FuncAnalysisState &funcState = getFuncAnalysisState(state);
+                                OneShotAnalysisState &state,
+                                FuncAnalysisState &funcState) {
   funcOp->walk([&](func::CallOp callOp) {
     func::FuncOp calledFunction = getCalledFunction(callOp);
     assert(calledFunction && "could not retrieved called func::FuncOp");
 
     // No equivalence info available for the called function.
     if (!funcState.equivalentFuncArgs.count(calledFunction))
       return WalkResult::skip();
 

 
 LogicalResult
 mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
                                      OneShotAnalysisState &state) {
   OneShotBufferizationOptions options =
       static_cast<const OneShotBufferizationOptions &>(state.getOptions());
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
-  FuncAnalysisState &funcState = getFuncAnalysisState(state);
+  FuncAnalysisState &funcState = getOrCreateFuncAnalysisState(state);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
-    equivalenceAnalysis(funcOp, aliasInfo, state);
+    equivalenceAnalysis(funcOp, aliasInfo, state, funcState);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state)))
       return failure();
 
     // Run some extra function analyses.
-    if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state)) ||
-        failed(funcOpBbArgReadWriteAnalysis(funcOp, state)))
+    if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state, funcState)) ||
+        failed(funcOpBbArgReadWriteAnalysis(funcOp, state, funcState)))
       return failure();
 
     // Mark op as fully analyzed.
     funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
   }
 
   return success();
 }