From 25886dd135b859a714e44441cbd16c94650b2db6 Mon Sep 17 00:00:00 2001
From: Mike Urbach <mike.urbach@sifive.com>
Date: Tue, 5 Mar 2024 13:08:03 -0800
Subject: [PATCH] [OM] Add C API and Python bindings for IntegerAttr to string.

Both the upstream MLIR IntegerAttr and OM IntegerAttr are backed by an
arbitrary precision integer. However, the upstream Python bindings
don't have any mechanism to return an integer larger than 64 bits back
to Python, even though Python ints are also arbitrary precision
integers.

To support this, we can handle this where we explicitly convert OM
IntegerAttrs to Python values. The simplest thing is to print a string
representation of the arbitrary precision integer, and parse that to a
Python int. This adds the necessary C API and Python binding for a "to
string" method, and uses it in the attribute_to_var function.

There are smarter ways we can handle the conversion, but the "to
string" API seems generally useful, so I'm using that in the
conversion for now.
---
 include/circt-c/Dialect/OM.h                  |  3 ++
 .../Bindings/Python/dialects/om.py            | 38 ++++++++++++++++++-
 lib/Bindings/Python/OMModule.cpp              |  8 +++-
 lib/Bindings/Python/support.py                |  2 +-
 lib/CAPI/Dialect/OM.cpp                       | 10 +++++
 5 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/include/circt-c/Dialect/OM.h b/include/circt-c/Dialect/OM.h
index 8149f05edd9..2248113d98c 100644
--- a/include/circt-c/Dialect/OM.h
+++ b/include/circt-c/Dialect/OM.h
@@ -236,6 +236,9 @@ MLIR_CAPI_EXPORTED MlirAttribute omIntegerAttrGetInt(MlirAttribute attr);
 /// Get an om::IntegerAttr from mlir::IntegerAttr.
 MLIR_CAPI_EXPORTED MlirAttribute omIntegerAttrGet(MlirAttribute attr);
 
+/// Get a string representation of an om::IntegerAttr.
+MLIR_CAPI_EXPORTED MlirStringRef omIntegerAttrToString(MlirAttribute attr);
+
 //===----------------------------------------------------------------------===//
 // ListAttr API
 //===----------------------------------------------------------------------===//
diff --git a/integration_test/Bindings/Python/dialects/om.py b/integration_test/Bindings/Python/dialects/om.py
index 891a2bac0ea..fda75f24cc7 100644
--- a/integration_test/Bindings/Python/dialects/om.py
+++ b/integration_test/Bindings/Python/dialects/om.py
@@ -3,7 +3,7 @@
 
 import circt
 from circt.dialects import om
-from circt.ir import Context, InsertionPoint, Location, Module
+from circt.ir import Context, InsertionPoint, Location, Module, IntegerAttr, IntegerType
 from circt.support import var_to_attribute
 
 from dataclasses import dataclass
@@ -259,3 +259,39 @@
 
 # CHECK: 3
 print(delayed.result)
+
+with Context() as ctx:
+  circt.register_dialects(ctx)
+
+  # Signless
+  int_attr1 = om.OMIntegerAttr.get(
+      IntegerAttr.get(IntegerType.get_signless(64), 42))
+  # CHECK: 42
+  print(str(int_attr1))
+
+  int_attr2 = om.OMIntegerAttr.get(
+      IntegerAttr.get(IntegerType.get_signless(64), -42))
+  # CHECK: 18446744073709551574
+  print(str(int_attr2))
+
+  # Signed
+  int_attr3 = om.OMIntegerAttr.get(
+      IntegerAttr.get(IntegerType.get_signed(64), 42))
+  # CHECK: 42
+  print(str(int_attr3))
+
+  int_attr4 = om.OMIntegerAttr.get(
+      IntegerAttr.get(IntegerType.get_signed(64), -42))
+  # CHECK: -42
+  print(str(int_attr4))
+
+  # Unsigned
+  int_attr5 = om.OMIntegerAttr.get(
+      IntegerAttr.get(IntegerType.get_unsigned(64), 42))
+  # CHECK: 42
+  print(str(int_attr5))
+
+  int_attr6 = om.OMIntegerAttr.get(
+      IntegerAttr.get(IntegerType.get_unsigned(64), -42))
+  # CHECK: 18446744073709551574
+  print(str(int_attr6))
diff --git a/lib/Bindings/Python/OMModule.cpp b/lib/Bindings/Python/OMModule.cpp
index 468c72019aa..a90226f5bef 100644
--- a/lib/Bindings/Python/OMModule.cpp
+++ b/lib/Bindings/Python/OMModule.cpp
@@ -476,8 +476,12 @@ void circt::python::populateDialectOMSubmodule(py::module &m) {
                        [](py::object cls, MlirAttribute intVal) {
                          return cls(omIntegerAttrGet(intVal));
                        })
-      .def_property_readonly("integer", [](MlirAttribute self) {
-        return omIntegerAttrGetInt(self);
+      .def_property_readonly(
+          "integer",
+          [](MlirAttribute self) { return omIntegerAttrGetInt(self); })
+      .def("__str__", [](MlirAttribute self) {
+        MlirStringRef str = omIntegerAttrToString(self);
+        return std::string(str.data, str.length);
       });
 
   // Add the OMListAttr definition
diff --git a/lib/Bindings/Python/support.py b/lib/Bindings/Python/support.py
index 50465a2abb1..0b906c8f9e3 100644
--- a/lib/Bindings/Python/support.py
+++ b/lib/Bindings/Python/support.py
@@ -195,7 +195,7 @@ def attribute_to_var(attr):
   except ValueError:
     pass
   try:
-    return attribute_to_var(om.OMIntegerAttr(attr).integer)
+    return int(str(om.OMIntegerAttr(attr)))
   except ValueError:
     pass
   try:
diff --git a/lib/CAPI/Dialect/OM.cpp b/lib/CAPI/Dialect/OM.cpp
index e0ee7fbd2ef..445ff3db6ca 100644
--- a/lib/CAPI/Dialect/OM.cpp
+++ b/lib/CAPI/Dialect/OM.cpp
@@ -381,6 +381,16 @@ MlirAttribute omIntegerAttrGet(MlirAttribute attr) {
       circt::om::IntegerAttr::get(integerAttr.getContext(), integerAttr));
 }
 
+/// Get a string representation of an om::IntegerAttr.
+MlirStringRef omIntegerAttrToString(MlirAttribute attr) {
+  mlir::IntegerAttr integerAttr =
+      cast<circt::om::IntegerAttr>(unwrap(attr)).getValue();
+  SmallVector<char> str;
+  integerAttr.getValue().toString(
+      str, /*Radix=*/10, /*Signed=*/integerAttr.getType().isSignedInteger());
+  return wrap(StringAttr::get(integerAttr.getContext(), str).getValue());
+}
+
 //===----------------------------------------------------------------------===//
 // ListAttr API.
 //===----------------------------------------------------------------------===//