From 7e4b8bd4e57a933acb732bd850f476a366de0e93 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            | 39 ++++++++++++++++++-
 lib/Bindings/Python/OMModule.cpp              |  8 +++-
 lib/Bindings/Python/support.py                |  2 +-
 lib/CAPI/Dialect/OM.cpp                       | 10 +++++
 5 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/include/circt-c/Dialect/OM.h b/include/circt-c/Dialect/OM.h
index 8149f05edd9c..2248113d98c1 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 891a2bac0ea8..3a3aa1a9ae22 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,40 @@
 
 # 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 468c72019aad..a90226f5befa 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 50465a2abb1c..0b906c8f9e3d 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 e0ee7fbd2ef5..445ff3db6caa 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.
 //===----------------------------------------------------------------------===//