Add Python bindings for OM unknown value (#9868)

Author: trmckay

include/circt-c/Dialect/OM.h

@@ -206,6 +206,18 @@ omEvaluatorValueIsAReference(OMEvaluatorValue evaluatorValue);
 MLIR_CAPI_EXPORTED OMEvaluatorValue
 omEvaluatorValueGetReferenceValue(OMEvaluatorValue evaluatorValue);
 
+/// Query if the EvaluatorValue is Unknown.
+MLIR_CAPI_EXPORTED bool
+omEvaluatorValueIsUnknown(OMEvaluatorValue evaluatorValue);
+
+/// Create an Unknown EvaluatorValue with the given type.
+MLIR_CAPI_EXPORTED OMEvaluatorValue omEvaluatorUnknownGet(MlirContext context,
+                                                          MlirType type);
+
+/// Get the type of an EvaluatorValue.
+MLIR_CAPI_EXPORTED MlirType
+omEvaluatorValueGetType(OMEvaluatorValue evaluatorValue);
+
 //===----------------------------------------------------------------------===//
 // ReferenceAttr API
 //===----------------------------------------------------------------------===//

integration_test/Bindings/Python/dialects/om.py

@@ -251,3 +251,62 @@
   list_type = Type.parse("!om.list<!om.any>")
   assert isinstance(list_type, om.ListType)
   assert isinstance(list_type.element_type, om.AnyType)
+
+  # Test partial evaluation with multiple inputs and outputs
+  module = Module.parse("""
+  module {
+    om.class @PartialEval(
+      %in0: !om.integer,
+      %in1: !om.integer,
+      %in2: !om.integer,
+      %in3: i1
+    ) -> (
+      out0: !om.integer,
+      out1: !om.integer,
+      out2: !om.integer,
+      out3: !om.integer,
+      out4: i1,
+      out5: i1
+    ) {
+      // out0: known constant
+      %c42 = om.constant #om.integer<42 : i32> : !om.integer
+
+      // out1: known input (in0)
+
+      // out2: unknown input (in1) propagated through arithmetic
+      %sum = om.integer.add %in1, %c42 : !om.integer
+
+      // out3: depends on unknown input (in2) and known input (in0)
+      %prod = om.integer.mul %in0, %in2 : !om.integer
+
+      // out4: known boolean constant
+      %true = om.constant true
+
+      // out5: unknown boolean input (in3)
+
+      om.class.fields %c42, %in0, %sum, %prod, %true, %in3 :
+        !om.integer, !om.integer, !om.integer, !om.integer, i1, i1
+    }
+  }
+  """)
+
+  evaluator = om.Evaluator(module)
+  om_integer_type = Type.parse("!om.integer")
+  i1_type = Type.parse("i1")
+
+  # Instantiate with mix of known and unknown inputs
+  obj = evaluator.instantiate("PartialEval", 100, om.Unknown(om_integer_type),
+                              om.Unknown(om_integer_type), om.Unknown(i1_type))
+
+  # CHECK: out0 (constant): 42
+  print(f"out0 (constant): {obj.out0}")
+  # CHECK: out1 (known input): 100
+  print(f"out1 (known input): {obj.out1}")
+  # CHECK: out2 (unknown input): Unknown(!om.integer)
+  print(f"out2 (unknown input): {obj.out2}")
+  # CHECK: out3 (depends on unknown): Unknown(!om.integer)
+  print(f"out3 (depends on unknown): {obj.out3}")
+  # CHECK: out4 (constant bool): True
+  print(f"out4 (constant bool): {obj.out4}")
+  # CHECK: out5 (unknown bool): Unknown(i1)
+  print(f"out5 (unknown bool): {obj.out5}")

lib/Bindings/Python/OMModule.cpp

@@ -13,6 +13,7 @@
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
@@ -29,6 +30,16 @@ struct Object;
 struct BasePath;
 struct Path;
 
+/// Represents a value that is not known because it is an unsupplied input, or
+/// derived from unsupplied inputs.
+struct Unknown {
+  Unknown(MlirType type) : type(type) {}
+  MlirType getType() const { return type; }
+
+private:
+  MlirType type;
+};
+
 /// These are the Python types that are represented by the different primitive
 /// OMEvaluatorValues as Attributes.
 using PythonPrimitive = std::variant<nb::int_, nb::float_, nb::str, nb::bool_,
@@ -41,7 +52,7 @@ using PythonPrimitive = std::variant<nb::int_, nb::float_, nb::str, nb::bool_,
 /// is tried first, then we can hit an assert inside the MLIR codebase.
 struct None {};
 using PythonValue =
-    std::variant<None, Object, List, BasePath, Path, PythonPrimitive>;
+    std::variant<None, Object, List, BasePath, Path, Unknown, PythonPrimitive>;
 
 /// Map an opaque OMEvaluatorValue into a python value.
 PythonValue omEvaluatorValueToPythonValue(OMEvaluatorValue result);
@@ -362,6 +373,9 @@ PythonValue omEvaluatorValueToPythonValue(OMEvaluatorValue result) {
   if (omEvaluatorValueIsAPath(result))
     return Path(result);
 
+  if (omEvaluatorValueIsUnknown(result))
+    return Unknown(omEvaluatorValueGetType(result));
+
   if (omEvaluatorValueIsAReference(result))
     return omEvaluatorValueToPythonValue(
         omEvaluatorValueGetReferenceValue(result));
@@ -385,6 +399,9 @@ OMEvaluatorValue pythonValueToOMEvaluatorValue(PythonValue result,
   if (auto *object = std::get_if<Object>(&result))
     return object->getValue();
 
+  if (auto *unknown = std::get_if<Unknown>(&result))
+    return omEvaluatorUnknownGet(ctx, unknown->getType());
+
   auto primitive = std::get<PythonPrimitive>(result);
   return omEvaluatorValueFromPrimitive(
       omPythonValueToPrimitive(primitive, ctx));
@@ -421,6 +438,19 @@ void circt::python::populateDialectOMSubmodule(nb::module_ &m) {
       .def(nb::init<Path>(), nb::arg("path"))
       .def("__str__", &Path::dunderStr);
 
+  // Add the Unknown sentinel class definition.
+  nb::class_<Unknown>(m, "Unknown")
+      .def(nb::init<MlirType>(), nb::arg("type"))
+      .def_prop_ro("type", &Unknown::getType)
+      .def("__repr__", [](const Unknown &u) {
+        PyPrintAccumulator printAccum;
+        printAccum.parts.append("Unknown(");
+        mlirTypePrint(u.getType(), printAccum.getCallback(),
+                      printAccum.getUserData());
+        printAccum.parts.append(")");
+        return printAccum.join();
+      });
+
   // Add the Object class definition.
   nb::class_<Object>(m, "Object")
       .def(nb::init<Object>(), nb::arg("object"))

lib/Bindings/Python/dialects/om.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from ._om_ops_gen import *
-from .._mlir_libs._circt._om import AnyType, Evaluator as BaseEvaluator, Object as BaseObject, List as BaseList, BasePath as BaseBasePath, BasePathType, Path, PathType, ClassType, ReferenceAttr, ListAttr, ListType, OMIntegerAttr
+from .._mlir_libs._circt._om import AnyType, Evaluator as BaseEvaluator, Object as BaseObject, List as BaseList, BasePath as BaseBasePath, BasePathType, Path, PathType, ClassType, ReferenceAttr, ListAttr, ListType, OMIntegerAttr, Unknown
 
 from ..ir import Attribute, Diagnostic, DiagnosticSeverity, Module, StringAttr, IntegerAttr, IntegerType
 from ..support import attribute_to_var, var_to_attribute
@@ -25,6 +25,9 @@ def wrap_mlir_object(value):
   if isinstance(value, (int, float, str, bool, tuple, list, dict)):
     return value
 
+  if isinstance(value, Unknown):
+    return value
+
   if isinstance(value, BaseList):
     return List(value)
 

lib/CAPI/Dialect/OM.cpp

@@ -321,6 +321,26 @@ omEvaluatorValueGetReferenceValue(OMEvaluatorValue evaluatorValue) {
   return wrap(result.value());
 }
 
+/// Query if the EvaluatorValue is an Unknown value.
+bool omEvaluatorValueIsUnknown(OMEvaluatorValue evaluatorValue) {
+  return unwrap(evaluatorValue)->isUnknown();
+}
+
+/// Create an Unknown EvaluatorValue.
+OMEvaluatorValue omEvaluatorUnknownGet(MlirContext context, MlirType type) {
+  auto *ctx = unwrap(context);
+  auto loc = UnknownLoc::get(ctx);
+  auto unknownValue =
+      evaluator::AttributeValue::get(unwrap(type), LocationAttr(loc));
+  unknownValue->markUnknown();
+  return wrap(unknownValue);
+}
+
+/// Get the type of an EvaluatorValue.
+MlirType omEvaluatorValueGetType(OMEvaluatorValue evaluatorValue) {
+  return wrap(unwrap(evaluatorValue)->getType());
+}
+
 //===----------------------------------------------------------------------===//
 // ReferenceAttr API.
 //===----------------------------------------------------------------------===//