Address review comments

llvm · Sep 19, 2023 · f0ce742 · f0ce742
1 parent c7a3ee1
commit f0ce742
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 12 deletions.
diff --git a/include/circt-c/Dialect/OM.h b/include/circt-c/Dialect/OM.h
@@ -193,9 +193,8 @@ MLIR_CAPI_EXPORTED bool omAttrIsAIntegerAttr(MlirAttribute attr);
 /// Given an om::IntegerAttr, return the mlir::IntegerAttr.
 MLIR_CAPI_EXPORTED MlirAttribute omIntegerAttrGetInt(MlirAttribute attr);
 
-/// If `attr` is an mlir::IntegerAttr, create and return the om::IntegerAttr,
-/// else return `attr` as is.
-MLIR_CAPI_EXPORTED MlirAttribute omCastIntAttrIfValid(MlirAttribute attr);
+/// Get an om::IntegerAttr from mlir::IntegerAttr.
+MLIR_CAPI_EXPORTED MlirAttribute omIntegerAttrGet(MlirAttribute attr);
 
 //===----------------------------------------------------------------------===//
 // ListAttr API

diff --git a/lib/Bindings/Python/OMModule.cpp b/lib/Bindings/Python/OMModule.cpp
@@ -26,6 +26,7 @@ struct List;
 struct Object;
 struct Tuple;
 struct Map;
+
 using PythonValue = std::variant<MlirAttribute, Object, List, Tuple, Map>;
 
 /// Map an opaque OMEvaluatorValue into a python value.
@@ -373,7 +374,7 @@ void circt::python::populateDialectOMSubmodule(py::module &m) {
   mlir_attribute_subclass(m, "OMIntegerAttr", omAttrIsAIntegerAttr)
       .def_classmethod("get",
                        [](py::object cls, MlirAttribute intVal) {
-                         return cls(omCastIntAttrIfValid(intVal));
+                         return cls(omIntegerAttrGet(intVal));
                        })
       .def_property_readonly("integer", [](MlirAttribute self) {
         return omIntegerAttrGetInt(self);

diff --git a/lib/Bindings/Python/dialects/om.py b/lib/Bindings/Python/dialects/om.py
@@ -38,6 +38,7 @@ def wrap_mlir_object(value):
   assert isinstance(value, BaseObject)
   return Object(value)
 
+
 def om_var_to_attribute(obj, none_on_fail: bool = False) -> ir.Attrbute:
   if isinstance(obj, int):
     return OMIntegerAttr.get(IntegerAttr.get(IntegerType.get_signless(64), obj))

diff --git a/lib/CAPI/Dialect/OM.cpp b/lib/CAPI/Dialect/OM.cpp
@@ -297,15 +297,13 @@ bool omAttrIsAIntegerAttr(MlirAttribute attr) {
 }
 
 MlirAttribute omIntegerAttrGetInt(MlirAttribute attr) {
-  return wrap(
-      (Attribute)unwrap(attr).cast<circt::om::IntegerAttr>().getValue());
+  return wrap(cast<circt::om::IntegerAttr>(unwrap(attr)).getValue());
 }
 
-MlirAttribute omCastIntAttrIfValid(MlirAttribute attr) {
-  if (auto integerAttr = llvm::dyn_cast<mlir::IntegerAttr>(unwrap(attr)))
-    return wrap(
-        circt::om::IntegerAttr::get(integerAttr.getContext(), integerAttr));
-  return attr;
+MlirAttribute omIntegerAttrGet(MlirAttribute attr) {
+  auto integerAttr = cast<mlir::IntegerAttr>(unwrap(attr));
+  return wrap(
+      circt::om::IntegerAttr::get(integerAttr.getContext(), integerAttr));
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/test/CAPI/om.c b/test/CAPI/om.c
@@ -57,7 +57,7 @@ void testEvaluator(MlirContext ctx) {
   // Test instantiation success.
 
   OMEvaluatorValue actualParam = omEvaluatorValueFromPrimitive(
-      omCastIntAttrIfValid(mlirIntegerAttrGet(mlirIntegerTypeGet(ctx, 8), 42)));
+      omIntegerAttrGet(mlirIntegerAttrGet(mlirIntegerTypeGet(ctx, 8), 42)));
 
   OMEvaluatorValue object =
       omEvaluatorInstantiate(evaluator, className, 1, &actualParam);