[OM] Add location info to EvaluatorValue (#6240)

Add the debug locations to the evaluator value, which will be used to generate info fields
 of the object model output json.

include/circt-c/Dialect/OM.h

@@ -127,6 +127,10 @@ omEvaluatorObjectGetFieldNames(OMEvaluatorValue object);
 MLIR_CAPI_EXPORTED MlirContext
 omEvaluatorValueGetContext(OMEvaluatorValue evaluatorValue);
 
+// Get Location from an EvaluatorValue.
+MLIR_CAPI_EXPORTED MlirLocation
+omEvaluatorValueGetLoc(OMEvaluatorValue evaluatorValue);
+
 // Query if the EvaluatorValue is null.
 MLIR_CAPI_EXPORTED bool omEvaluatorValueIsNull(OMEvaluatorValue evaluatorValue);
 

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

@@ -17,6 +17,7 @@
 #include "circt/Support/LLVM.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LogicalResult.h"
@@ -46,7 +47,8 @@ using ObjectFields = SmallDenseMap<StringAttr, EvaluatorValuePtr>;
 struct EvaluatorValue : std::enable_shared_from_this<EvaluatorValue> {
   // Implement LLVM RTTI.
   enum class Kind { Attr, Object, List, Tuple, Map, Reference };
-  EvaluatorValue(MLIRContext *ctx, Kind kind) : kind(kind), ctx(ctx) {}
+  EvaluatorValue(MLIRContext *ctx, Kind kind, Location loc)
+      : kind(kind), ctx(ctx), loc(loc) {}
   Kind getKind() const { return kind; }
   MLIRContext *getContext() const { return ctx; }
 
@@ -63,9 +65,20 @@ struct EvaluatorValue : std::enable_shared_from_this<EvaluatorValue> {
   // Finalize the evaluator value. Strip intermidiate reference values.
   LogicalResult finalize();
 
+  // Return the Location associated with the Value.
+  Location getLoc() const { return loc; }
+  // Set the Location associated with the Value.
+  void setLoc(Location l) { loc = l; }
+  // Set the Location, if it is unknown.
+  void setLocIfUnknown(Location l) {
+    if (isa<UnknownLoc>(loc))
+      loc = l;
+  }
+
 private:
   const Kind kind;
   MLIRContext *ctx;
+  Location loc;
   bool fullyEvaluated = false;
   bool finalized = false;
 };
@@ -74,8 +87,8 @@ struct EvaluatorValue : std::enable_shared_from_this<EvaluatorValue> {
 /// 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),
+  ReferenceValue(Type type, Location loc)
+      : EvaluatorValue(type.getContext(), Kind::Reference, loc), value(nullptr),
         type(type) {}
 
   // Implement LLVM RTTI.
@@ -114,10 +127,11 @@ struct ReferenceValue : EvaluatorValue {
 /// Values which can be directly representable by MLIR attributes.
 struct AttributeValue : EvaluatorValue {
   AttributeValue(Attribute attr)
-      : EvaluatorValue(attr.getContext(), Kind::Attr), attr(attr) {
+      : AttributeValue(attr, mlir::UnknownLoc::get(attr.getContext())) {}
+  AttributeValue(Attribute attr, Location loc)
+      : EvaluatorValue(attr.getContext(), Kind::Attr, loc), attr(attr) {
     markFullyEvaluated();
   }
-
   Attribute getAttr() const { return attr; }
   template <typename AttrTy>
   AttrTy getAs() const {
@@ -151,8 +165,9 @@ static inline LogicalResult finalizeEvaluatorValue(EvaluatorValuePtr &value) {
 
 /// 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),
+  ListValue(om::ListType type, SmallVector<EvaluatorValuePtr> elements,
+            Location loc)
+      : EvaluatorValue(type.getContext(), Kind::List, loc), type(type),
         elements(std::move(elements)) {
     markFullyEvaluated();
   }
@@ -166,8 +181,8 @@ struct ListValue : EvaluatorValue {
   LogicalResult finalizeImpl();
 
   // Partially evaluated value.
-  ListValue(om::ListType type)
-      : EvaluatorValue(type.getContext(), Kind::List), type(type) {}
+  ListValue(om::ListType type, Location loc)
+      : EvaluatorValue(type.getContext(), Kind::List, loc), type(type) {}
 
   const auto &getElements() const { return elements; }
 
@@ -186,15 +201,16 @@ struct ListValue : EvaluatorValue {
 
 /// A Map value.
 struct MapValue : EvaluatorValue {
-  MapValue(om::MapType type, DenseMap<Attribute, EvaluatorValuePtr> elements)
-      : EvaluatorValue(type.getContext(), Kind::Map), type(type),
+  MapValue(om::MapType type, DenseMap<Attribute, EvaluatorValuePtr> elements,
+           Location loc)
+      : EvaluatorValue(type.getContext(), Kind::Map, loc), type(type),
         elements(std::move(elements)) {
     markFullyEvaluated();
   }
 
   // Partially evaluated value.
-  MapValue(om::MapType type)
-      : EvaluatorValue(type.getContext(), Kind::Map), type(type) {}
+  MapValue(om::MapType type, Location loc)
+      : EvaluatorValue(type.getContext(), Kind::Map, loc), type(type) {}
 
   const auto &getElements() const { return elements; }
   void setElements(DenseMap<Attribute, EvaluatorValuePtr> newElements) {
@@ -223,15 +239,15 @@ struct MapValue : EvaluatorValue {
 
 /// A composite Object, which has a type and fields.
 struct ObjectValue : EvaluatorValue {
-  ObjectValue(om::ClassOp cls, ObjectFields fields)
-      : EvaluatorValue(cls.getContext(), Kind::Object), cls(cls),
+  ObjectValue(om::ClassOp cls, ObjectFields fields, Location loc)
+      : EvaluatorValue(cls.getContext(), Kind::Object, loc), cls(cls),
         fields(std::move(fields)) {
     markFullyEvaluated();
   }
 
   // Partially evaluated value.
-  ObjectValue(om::ClassOp cls)
-      : EvaluatorValue(cls.getContext(), Kind::Object), cls(cls) {}
+  ObjectValue(om::ClassOp cls, Location loc)
+      : EvaluatorValue(cls.getContext(), Kind::Object, loc), cls(cls) {}
 
   om::ClassOp getClassOp() const { return cls; }
   const auto &getFields() const { return fields; }
@@ -275,15 +291,15 @@ struct ObjectValue : EvaluatorValue {
 /// Tuple values.
 struct TupleValue : EvaluatorValue {
   using TupleElements = llvm::SmallVector<EvaluatorValuePtr>;
-  TupleValue(TupleType type, TupleElements tupleElements)
-      : EvaluatorValue(type.getContext(), Kind::Tuple), type(type),
+  TupleValue(TupleType type, TupleElements tupleElements, Location loc)
+      : EvaluatorValue(type.getContext(), Kind::Tuple, loc), type(type),
         elements(std::move(tupleElements)) {
     markFullyEvaluated();
   }
 
   // Partially evaluated value.
-  TupleValue(TupleType type)
-      : EvaluatorValue(type.getContext(), Kind::Tuple), type(type) {}
+  TupleValue(TupleType type, Location loc)
+      : EvaluatorValue(type.getContext(), Kind::Tuple, loc), type(type) {}
 
   void setElements(TupleElements newElements) {
     elements = std::move(newElements);
@@ -334,7 +350,8 @@ struct Evaluator {
   /// Get the Module this Evaluator is built from.
   mlir::ModuleOp getModule();
 
-  FailureOr<evaluator::EvaluatorValuePtr> getPartiallyEvaluatedValue(Type type);
+  FailureOr<evaluator::EvaluatorValuePtr>
+  getPartiallyEvaluatedValue(Type type, Location loc);
 
   using ActualParameters =
       SmallVectorImpl<std::shared_ptr<evaluator::EvaluatorValue>> *;
@@ -351,42 +368,48 @@ struct Evaluator {
     return val && val->isFullyEvaluated();
   }
 
-  FailureOr<EvaluatorValuePtr> getOrCreateValue(Value value,
-                                                ActualParameters actualParams);
+  FailureOr<EvaluatorValuePtr>
+  getOrCreateValue(Value value, ActualParameters actualParams, Location loc);
   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,
-                                             ActualParameters actualParams);
+  FailureOr<EvaluatorValuePtr>
+  evaluateValue(Value value, ActualParameters actualParams, Location loc);
 
   /// Evaluator dispatch functions for the small expression grammar.
   FailureOr<EvaluatorValuePtr> evaluateParameter(BlockArgument formalParam,
-                                                 ActualParameters actualParams);
+                                                 ActualParameters actualParams,
+                                                 Location loc);
 
-  FailureOr<EvaluatorValuePtr> evaluateConstant(ConstantOp op,
-                                                ActualParameters actualParams);
+  FailureOr<EvaluatorValuePtr>
+  evaluateConstant(ConstantOp op, ActualParameters actualParams, Location loc);
   /// Instantiate an Object with its class name and actual parameters.
   FailureOr<EvaluatorValuePtr>
   evaluateObjectInstance(StringAttr className, ActualParameters actualParams,
-                         ObjectKey instanceObjectKey = {});
+                         Location loc, ObjectKey instanceKey = {});
   FailureOr<EvaluatorValuePtr>
   evaluateObjectInstance(ObjectOp op, ActualParameters actualParams);
   FailureOr<EvaluatorValuePtr>
-  evaluateObjectField(ObjectFieldOp op, ActualParameters actualParams);
+  evaluateObjectField(ObjectFieldOp op, ActualParameters actualParams,
+                      Location loc);
+  FailureOr<EvaluatorValuePtr> evaluateListCreate(ListCreateOp op,
+                                                  ActualParameters actualParams,
+                                                  Location loc);
   FailureOr<EvaluatorValuePtr>
-  evaluateListCreate(ListCreateOp op, ActualParameters actualParams);
+  evaluateTupleCreate(TupleCreateOp op, ActualParameters actualParams,
+                      Location loc);
   FailureOr<EvaluatorValuePtr>
-  evaluateTupleCreate(TupleCreateOp op, ActualParameters actualParams);
-  FailureOr<EvaluatorValuePtr> evaluateTupleGet(TupleGetOp op,
-                                                ActualParameters actualParams);
+  evaluateTupleGet(TupleGetOp op, ActualParameters actualParams, Location loc);
   FailureOr<evaluator::EvaluatorValuePtr>
-  evaluateMapCreate(MapCreateOp op, ActualParameters actualParams);
+  evaluateMapCreate(MapCreateOp op, ActualParameters actualParams,
+                    Location loc);
 
   FailureOr<ActualParameters>
-  createParametersFromOperands(ValueRange range, ActualParameters actualParams);
+  createParametersFromOperands(ValueRange range, ActualParameters actualParams,
+                               Location loc);
 
   /// The symbol table for the IR module the Evaluator was constructed with.
   /// Used to look up class definitions.

integration_test/Bindings/Python/dialects/om.py

@@ -122,16 +122,29 @@
 
 # CHECK: 42
 print(obj.field)
+
+# location of the om.class.field @field
+# CHECK: loc("-":28:7)
+print(obj.get_field_loc("field"))
+
 # CHECK: 14
 print(obj.child.foo)
+# CHECK: loc("-":64:7)
+print(obj.child.get_field_loc("foo"))
 # CHECK: ('Root', 'x')
 print(obj.reference)
-# CHECK: 14
 (fst, snd) = obj.tuple
+# CHECK: 14
 print(snd)
 
+# CHECK: loc("-":43:7)
+print(obj.get_field_loc("tuple"))
+
 # CHECK: path
 print(obj.path)
+# location of om.class.field @path, %path : !om.path
+# CHECK: loc("-":35:7)
+print(obj.get_field_loc("path"))
 
 try:
   print(obj.tuple[3])
@@ -140,10 +153,17 @@
   print(e)
 
 for (name, field) in obj:
+  # location from om.class.field @child, %0 : !om.class.type<@Child>
   # CHECK: name: child, field: <circt.dialects.om.Object object
+  # CHECK-SAME: loc: loc("-":32:7)
+  # location from om.class.field @field, %param : !om.integer
   # CHECK: name: field, field: 42
+  # CHECK-SAME: loc: loc("-":28:7)
+  # location from om.class.field @reference, %sym : !om.ref
   # CHECK: name: reference, field: ('Root', 'x')
-  print(f"name: {name}, field: {field}")
+  # CHECK-SAME: loc: loc("-":37:7)
+  loc = obj.get_field_loc(name)
+  print(f"name: {name}, field: {field}, loc: {loc}")
 
 # CHECK: ['X', 'Y']
 print(obj.list)

lib/Bindings/Python/OMModule.cpp

@@ -107,6 +107,19 @@ struct Object {
   /// Get the Type from an Object, which will be a ClassType.
   MlirType getType() { return omEvaluatorObjectGetType(value); }
 
+  // Get the field location info.
+  MlirLocation getFieldLoc(const std::string &name) {
+    // Wrap the requested field name in an attribute.
+    MlirContext context = mlirTypeGetContext(omEvaluatorObjectGetType(value));
+    MlirStringRef cName = mlirStringRefCreateFromCString(name.c_str());
+    MlirAttribute nameAttr = mlirStringAttrGet(context, cName);
+
+    // Get the field's ObjectValue via the CAPI.
+    OMEvaluatorValue result = omEvaluatorObjectGetField(value, nameAttr);
+
+    return omEvaluatorValueGetLoc(result);
+  }
+
   // Get a field from the Object, using pybind's support for variant to return a
   // Python object that is either an Object or Attribute.
   PythonValue getField(const std::string &name) {
@@ -365,6 +378,8 @@ void circt::python::populateDialectOMSubmodule(py::module &m) {
       .def(py::init<Object>(), py::arg("object"))
       .def("__getattr__", &Object::getField, "Get a field from an Object",
            py::arg("name"))
+      .def("get_field_loc", &Object::getFieldLoc,
+           "Get the location of a field from an Object", py::arg("name"))
       .def_property_readonly("field_names", &Object::getFieldNames,
                              "Get field names from an Object")
       .def_property_readonly("type", &Object::getType, "The Type of the Object")

lib/Bindings/Python/dialects/om.py

@@ -133,6 +133,11 @@ def __getattr__(self, name: str):
     field = super().__getattr__(name)
     return wrap_mlir_object(field)
 
+  def get_field_loc(self, name: str):
+    # Call the base method to get the loc.
+    loc = super().get_field_loc(name)
+    return loc
+
   # Support iterating over an Object by yielding its fields.
   def __iter__(self):
     for name in self.field_names:

lib/CAPI/Dialect/OM.cpp

@@ -16,6 +16,7 @@
 #include "circt/Dialect/OM/OMDialect.h"
 #include "mlir/CAPI/Registration.h"
 #include "mlir/CAPI/Wrap.h"
+#include "mlir/IR/Location.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/Support/Casting.h"
 
@@ -191,6 +192,11 @@ MlirContext omEvaluatorValueGetContext(OMEvaluatorValue evaluatorValue) {
   return wrap(unwrap(evaluatorValue)->getContext());
 }
 
+// Get location from an EvaluatorValue.
+MlirLocation omEvaluatorValueGetLoc(OMEvaluatorValue evaluatorValue) {
+  return wrap(unwrap(evaluatorValue)->getLoc());
+}
+
 // Query if the EvaluatorValue is null.
 bool omEvaluatorValueIsNull(OMEvaluatorValue evaluatorValue) {
   // Check if the pointer is null.