[OM] Evaluator: Support graph regions (#6249)

This commit adds support for graph regions for evaluator.

`ReferenceValue`, a new subclass of `EvaluatorValue`, is added to behave as pointers. `ReferenceValue` can be used as alias to different values and is created for `class.object.field` operation because `class.object.field` can access fields across class hierarchies and the fields might not be initilized yet. `RefenceValue` is not exposed to outside of evaluator implemenation. `EvaluatorValue::finalize` shrinks intermidiate `RefenceValue` in the evaluator value.

Evaluation algorithm is changed to worklist-based iteration. `instantiae` method first traverses the whole IR including sub-class, and create partially evaluaed values for all values and add these values to the worklist. After that we evaluate values until there is no partially evaluaed value.

Fix #5834

Author: uenoku

include/circt/Dialect/OM/Evaluator/Evaluator.h

@@ -20,6 +20,10 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/SmallPtrSet.h"
+
+#include <queue>
+#include <utility>
 
 namespace circt {
 namespace om {
@@ -36,25 +40,84 @@ using EvaluatorValuePtr = std::shared_ptr<EvaluatorValue>;
 using ObjectFields = SmallDenseMap<StringAttr, EvaluatorValuePtr>;
 
 /// Base class for evaluator runtime values.
-/// Enables the shared_from_this functionality so Object pointers can be passed
-/// through the CAPI and unwrapped back into C++ smart pointers with the
-/// appropriate reference count.
+/// Enables the shared_from_this functionality so Evaluator Value pointers can
+/// be passed through the CAPI and unwrapped back into C++ smart pointers with
+/// the appropriate reference count.
 struct EvaluatorValue : std::enable_shared_from_this<EvaluatorValue> {
   // Implement LLVM RTTI.
-  enum class Kind { Attr, Object, List, Tuple, Map };
-  EvaluatorValue(MLIRContext *ctx, Kind kind) : kind(kind) {}
+  enum class Kind { Attr, Object, List, Tuple, Map, Reference };
+  EvaluatorValue(MLIRContext *ctx, Kind kind) : kind(kind), ctx(ctx) {}
   Kind getKind() const { return kind; }
   MLIRContext *getContext() const { return ctx; }
 
+  // Return true the value is fully evaluated.
+  bool isFullyEvaluated() const { return fullyEvaluated; }
+  void markFullyEvaluated() {
+    assert(!fullyEvaluated && "should not mark twice");
+    fullyEvaluated = true;
+  }
+
+  // Return a MLIR type which the value represents.
+  Type getType() const;
+
+  // Finalize the evaluator value. Strip intermidiate reference values.
+  LogicalResult finalize();
+
 private:
   const Kind kind;
   MLIRContext *ctx;
+  bool fullyEvaluated = false;
+  bool finalized = false;
+};
+
+/// Values which can be used as pointers to different values.
+/// ReferenceValue is replaced with its element and erased at the end of
+/// evaluation.
+struct ReferenceValue : EvaluatorValue {
+  ReferenceValue(Type type)
+      : EvaluatorValue(type.getContext(), Kind::Reference), value(nullptr),
+        type(type) {}
+
+  // Implement LLVM RTTI.
+  static bool classof(const EvaluatorValue *e) {
+    return e->getKind() == Kind::Reference;
+  }
+
+  Type getValueType() const { return type; }
+  EvaluatorValuePtr getValue() const { return value; }
+  void setValue(EvaluatorValuePtr newValue) {
+    value = std::move(newValue);
+    markFullyEvaluated();
+  }
+
+  // Finalize the value.
+  LogicalResult finalizeImpl();
+
+  // Return the first non-reference value that is reachable from the reference.
+  FailureOr<EvaluatorValuePtr> getStrippedValue() const {
+    llvm::SmallPtrSet<ReferenceValue *, 4> visited;
+    auto currentValue = value;
+    while (auto *v = dyn_cast<ReferenceValue>(currentValue.get())) {
+      // Detect a cycle.
+      if (!visited.insert(v).second)
+        return failure();
+      currentValue = v->getValue();
+    }
+    return success(currentValue);
+  }
+
+private:
+  EvaluatorValuePtr value;
+  Type type;
 };
 
 /// Values which can be directly representable by MLIR attributes.
 struct AttributeValue : EvaluatorValue {
   AttributeValue(Attribute attr)
-      : EvaluatorValue(attr.getContext(), Kind::Attr), attr(attr) {}
+      : EvaluatorValue(attr.getContext(), Kind::Attr), attr(attr) {
+    markFullyEvaluated();
+  }
+
   Attribute getAttr() const { return attr; }
   template <typename AttrTy>
   AttrTy getAs() const {
@@ -64,20 +127,52 @@ struct AttributeValue : EvaluatorValue {
     return e->getKind() == Kind::Attr;
   }
 
+  // Finalize the value.
+  LogicalResult finalizeImpl() { return success(); }
+
+  Type getType() const { return attr.cast<TypedAttr>().getType(); }
+
 private:
-  Attribute attr;
+  Attribute attr = {};
 };
 
+// This perform finalization to `value`.
+static inline LogicalResult finalizeEvaluatorValue(EvaluatorValuePtr &value) {
+  if (failed(value->finalize()))
+    return failure();
+  if (auto *ref = llvm::dyn_cast<ReferenceValue>(value.get())) {
+    auto v = ref->getStrippedValue();
+    if (failed(v))
+      return v;
+    value = v.value();
+  }
+  return success();
+}
+
 /// A List which contains variadic length of elements with the same type.
 struct ListValue : EvaluatorValue {
   ListValue(om::ListType type, SmallVector<EvaluatorValuePtr> elements)
       : EvaluatorValue(type.getContext(), Kind::List), type(type),
-        elements(std::move(elements)) {}
+        elements(std::move(elements)) {
+    markFullyEvaluated();
+  }
+
+  void setElements(SmallVector<EvaluatorValuePtr> newElements) {
+    elements = std::move(newElements);
+    markFullyEvaluated();
+  }
+
+  // Finalize the value.
+  LogicalResult finalizeImpl();
+
+  // Partially evaluated value.
+  ListValue(om::ListType type)
+      : EvaluatorValue(type.getContext(), Kind::List), type(type) {}
 
   const auto &getElements() const { return elements; }
 
   /// Return the type of the value, which is a ListType.
-  om::ListType getType() const { return type; }
+  om::ListType getListType() const { return type; }
 
   /// Implement LLVM RTTI.
   static bool classof(const EvaluatorValue *e) {
@@ -93,12 +188,25 @@ struct ListValue : EvaluatorValue {
 struct MapValue : EvaluatorValue {
   MapValue(om::MapType type, DenseMap<Attribute, EvaluatorValuePtr> elements)
       : EvaluatorValue(type.getContext(), Kind::Map), type(type),
-        elements(std::move(elements)) {}
+        elements(std::move(elements)) {
+    markFullyEvaluated();
+  }
+
+  // Partially evaluated value.
+  MapValue(om::MapType type)
+      : EvaluatorValue(type.getContext(), Kind::Map), type(type) {}
 
   const auto &getElements() const { return elements; }
+  void setElements(DenseMap<Attribute, EvaluatorValuePtr> newElements) {
+    elements = std::move(newElements);
+    markFullyEvaluated();
+  }
+
+  // Finalize the evaluator value.
+  LogicalResult finalizeImpl();
 
   /// Return the type of the value, which is a MapType.
-  om::MapType getType() const { return type; }
+  om::MapType getMapType() const { return type; }
 
   /// Return an array of keys in the ascending order.
   ArrayAttr getKeys();
@@ -117,28 +225,48 @@ struct MapValue : EvaluatorValue {
 struct ObjectValue : EvaluatorValue {
   ObjectValue(om::ClassOp cls, ObjectFields fields)
       : EvaluatorValue(cls.getContext(), Kind::Object), cls(cls),
-        fields(std::move(fields)) {}
+        fields(std::move(fields)) {
+    markFullyEvaluated();
+  }
+
+  // Partially evaluated value.
+  ObjectValue(om::ClassOp cls)
+      : EvaluatorValue(cls.getContext(), Kind::Object), cls(cls) {}
+
   om::ClassOp getClassOp() const { return cls; }
   const auto &getFields() const { return fields; }
 
+  void setFields(llvm::SmallDenseMap<StringAttr, EvaluatorValuePtr> newFields) {
+    fields = std::move(newFields);
+    markFullyEvaluated();
+  }
+
   /// Return the type of the value, which is a ClassType.
-  om::ClassType getType() const {
+  om::ClassType getObjectType() const {
     auto clsConst = const_cast<ClassOp &>(cls);
     return ClassType::get(clsConst.getContext(),
                           FlatSymbolRefAttr::get(clsConst.getNameAttr()));
   }
 
+  Type getType() const { return getObjectType(); }
+
   /// Implement LLVM RTTI.
   static bool classof(const EvaluatorValue *e) {
     return e->getKind() == Kind::Object;
   }
 
   /// Get a field of the Object by name.
   FailureOr<EvaluatorValuePtr> getField(StringAttr field);
+  FailureOr<EvaluatorValuePtr> getField(StringRef field) {
+    return getField(StringAttr::get(getContext(), field));
+  }
 
   /// Get all the field names of the Object.
   ArrayAttr getFieldNames();
 
+  // Finalize the evaluator value.
+  LogicalResult finalizeImpl();
+
 private:
   om::ClassOp cls;
   llvm::SmallDenseMap<StringAttr, EvaluatorValuePtr> fields;
@@ -149,15 +277,33 @@ struct TupleValue : EvaluatorValue {
   using TupleElements = llvm::SmallVector<EvaluatorValuePtr>;
   TupleValue(TupleType type, TupleElements tupleElements)
       : EvaluatorValue(type.getContext(), Kind::Tuple), type(type),
-        elements(std::move(tupleElements)) {}
+        elements(std::move(tupleElements)) {
+    markFullyEvaluated();
+  }
+
+  // Partially evaluated value.
+  TupleValue(TupleType type)
+      : EvaluatorValue(type.getContext(), Kind::Tuple), type(type) {}
+
+  void setElements(TupleElements newElements) {
+    elements = std::move(newElements);
+    markFullyEvaluated();
+  }
+
+  LogicalResult finalizeImpl() {
+    for (auto &&value : elements)
+      if (failed(finalizeEvaluatorValue(value)))
+        return failure();
 
+    return success();
+  }
   /// Implement LLVM RTTI.
   static bool classof(const EvaluatorValue *e) {
     return e->getKind() == Kind::Tuple;
   }
 
   /// Return the type of the value, which is a TupleType.
-  TupleType getType() const { return type; }
+  TupleType getTupleType() const { return type; }
 
   const TupleElements &getElements() const { return elements; }
 
@@ -182,49 +328,81 @@ struct Evaluator {
   Evaluator(ModuleOp mod);
 
   /// Instantiate an Object with its class name and actual parameters.
-  FailureOr<std::shared_ptr<Object>>
+  FailureOr<evaluator::EvaluatorValuePtr>
   instantiate(StringAttr className, ArrayRef<EvaluatorValuePtr> actualParams);
 
   /// Get the Module this Evaluator is built from.
   mlir::ModuleOp getModule();
 
+  FailureOr<evaluator::EvaluatorValuePtr> getPartiallyEvaluatedValue(Type type);
+
+  using ActualParameters =
+      SmallVectorImpl<std::shared_ptr<evaluator::EvaluatorValue>> *;
+
+  using ObjectKey = std::pair<Value, ActualParameters>;
+
 private:
+  bool isFullyEvaluated(Value value, ActualParameters key) {
+    return isFullyEvaluated({value, key});
+  }
+
+  bool isFullyEvaluated(ObjectKey key) {
+    auto val = objects.lookup(key);
+    return val && val->isFullyEvaluated();
+  }
+
+  FailureOr<EvaluatorValuePtr> getOrCreateValue(Value value,
+                                                ActualParameters actualParams);
+  FailureOr<EvaluatorValuePtr>
+  allocateObjectInstance(StringAttr clasName, ActualParameters actualParams);
+
   /// Evaluate a Value in a Class body according to the small expression grammar
   /// described in the rationale document. The actual parameters are the values
   /// supplied at the current instantiation of the Class being evaluated.
-  FailureOr<EvaluatorValuePtr>
-  evaluateValue(Value value, ArrayRef<EvaluatorValuePtr> actualParams);
+  FailureOr<EvaluatorValuePtr> evaluateValue(Value value,
+                                             ActualParameters actualParams);
 
   /// Evaluator dispatch functions for the small expression grammar.
-  FailureOr<EvaluatorValuePtr>
-  evaluateParameter(BlockArgument formalParam,
-                    ArrayRef<EvaluatorValuePtr> actualParams);
+  FailureOr<EvaluatorValuePtr> evaluateParameter(BlockArgument formalParam,
+                                                 ActualParameters actualParams);
 
+  FailureOr<EvaluatorValuePtr> evaluateConstant(ConstantOp op,
+                                                ActualParameters actualParams);
+  /// Instantiate an Object with its class name and actual parameters.
   FailureOr<EvaluatorValuePtr>
-  evaluateConstant(ConstantOp op, ArrayRef<EvaluatorValuePtr> actualParams);
-  FailureOr<EvaluatorValuePtr>
-  evaluateObjectInstance(ObjectOp op, ArrayRef<EvaluatorValuePtr> actualParams);
+  evaluateObjectInstance(StringAttr className, ActualParameters actualParams,
+                         ObjectKey instanceObjectKey = {});
   FailureOr<EvaluatorValuePtr>
-  evaluateObjectField(ObjectFieldOp op,
-                      ArrayRef<EvaluatorValuePtr> actualParams);
+  evaluateObjectInstance(ObjectOp op, ActualParameters actualParams);
   FailureOr<EvaluatorValuePtr>
-  evaluateListCreate(ListCreateOp op, ArrayRef<EvaluatorValuePtr> actualParams);
+  evaluateObjectField(ObjectFieldOp op, ActualParameters actualParams);
   FailureOr<EvaluatorValuePtr>
-  evaluateTupleCreate(TupleCreateOp op,
-                      ArrayRef<EvaluatorValuePtr> actualParams);
+  evaluateListCreate(ListCreateOp op, ActualParameters actualParams);
   FailureOr<EvaluatorValuePtr>
-  evaluateTupleGet(TupleGetOp op, ArrayRef<EvaluatorValuePtr> actualParams);
+  evaluateTupleCreate(TupleCreateOp op, ActualParameters actualParams);
+  FailureOr<EvaluatorValuePtr> evaluateTupleGet(TupleGetOp op,
+                                                ActualParameters actualParams);
   FailureOr<evaluator::EvaluatorValuePtr>
-  evaluateMapCreate(MapCreateOp op,
-                    ArrayRef<evaluator::EvaluatorValuePtr> actualParams);
+  evaluateMapCreate(MapCreateOp op, ActualParameters actualParams);
+
+  FailureOr<ActualParameters>
+  createParametersFromOperands(ValueRange range, ActualParameters actualParams);
 
   /// The symbol table for the IR module the Evaluator was constructed with.
   /// Used to look up class definitions.
   SymbolTable symbolTable;
 
-  /// Object storage. Currently used for memoizing calls to
-  /// evaluateObjectInstance. Further refinement is expected.
-  DenseMap<Value, std::shared_ptr<evaluator::EvaluatorValue>> objects;
+  /// This uniquely stores vectors that represent parameters.
+  SmallVector<
+      std::unique_ptr<SmallVector<std::shared_ptr<evaluator::EvaluatorValue>>>>
+      actualParametersBuffers;
+
+  /// A worklist that tracks values which needs to be fully evaluated.
+  std::queue<ObjectKey> worklist;
+
+  /// Evaluator value storage. Return an evaluator value for the given
+  /// instantiation context (a pair of Value and parameters).
+  DenseMap<ObjectKey, std::shared_ptr<evaluator::EvaluatorValue>> objects;
 };
 
 /// Helper to enable printing objects in Diagnostics.

integration_test/Bindings/Python/dialects/om.py

@@ -185,3 +185,7 @@
   if isinstance(data, om.Object):
     object_dict[data] = field_name
 assert len(object_dict) == 2
+
+obj = evaluator.instantiate("Test", 41)
+# CHECK: 41
+print(obj.field)

lib/CAPI/Dialect/OM.cpp

@@ -108,8 +108,7 @@ OMEvaluatorValue omEvaluatorInstantiate(OMEvaluator evaluator,
     cppActualParams.push_back(unwrap(actualParams[i]));
 
   // Invoke the Evaluator to instantiate the Object.
-  FailureOr<std::shared_ptr<evaluator::ObjectValue>> result =
-      cppEvaluator->instantiate(cppClassName, cppActualParams);
+  auto result = cppEvaluator->instantiate(cppClassName, cppActualParams);
 
   // If instantiation failed, return a null Object. A Diagnostic will be emitted
   // in this case.