Repository: qir-alliance/pyqir (language: python)

Issue:
Add custom attributes to functions
In the Python API, add the ability to add custom attribute metadata when defining functions.


Discussion / hints:

Fix patch:
diff --git a/pyqir/pyqir/__init__.py b/pyqir/pyqir/__init__.py
index 2f555c07..22e3d471 100644
--- a/pyqir/pyqir/__init__.py
+++ b/pyqir/pyqir/__init__.py
@@ -34,10 +34,10 @@
     Switch,
     Type,
     Value,
+    add_string_attribute,
     const,
     dynamic_qubit_management,
     dynamic_result_management,
-    entry_point,
     extract_byte_string,
     global_byte_string,
     is_entry_point,
@@ -57,6 +57,7 @@
     result_type,
 )
 from pyqir._simple import SimpleModule
+from pyqir._entry_point import entry_point
 from pyqir._basicqis import BasicQisBuilder
 
 __all__ = [
@@ -94,6 +95,7 @@
     "Switch",
     "Type",
     "Value",
+    "add_string_attribute",
     "const",
     "dynamic_qubit_management",
     "dynamic_result_management",
diff --git a/pyqir/pyqir/_entry_point.py b/pyqir/pyqir/_entry_point.py
new file mode 100644
index 00000000..12d57e85
--- /dev/null
+++ b/pyqir/pyqir/_entry_point.py
@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import List, Optional
+
+import pyqir
+from pyqir import Module, Function, Context, add_string_attribute, Linkage, FunctionType
+
+
+def entry_point(
+    module: Module,
+    name: str,
+    required_num_qubits: int,
+    required_num_results: int,
+    qir_profiles: str = "custom",
+    output_labeling_schema: Optional[str] = None,
+) -> Function:
+    """
+    Creates an entry point.
+
+    :param Module module: The parent module.
+    :param str name: The entry point name.
+    :param int required_num_qubits: The number of qubits required by the entry point.
+    :param int required_num_results: The number of results required by the entry point.
+    :param Optional[str] qir_profiles: Value identifying the profile the entry point has been compiled for. Use base_profile when QIR is compliant.
+    :param Optional[str] output_labeling_schema: An arbitrary string value that identifies the schema used by a compiler frontend that produced the IR to label the recorded output
+    :returns: An entry point.
+    """
+    void = pyqir.Type.void(module.context)
+    function = Function(FunctionType(void, []), Linkage.EXTERNAL, name, module)
+    add_string_attribute(function, "entry_point")
+    add_string_attribute(function, "num_required_qubits", str(required_num_qubits))
+    add_string_attribute(function, "num_required_results", str(required_num_results))
+    add_string_attribute(function, "qir_profiles", qir_profiles)
+
+    add_string_attribute(
+        function,
+        "output_labeling_schema",
+        output_labeling_schema,
+    )
+
+    return function
diff --git a/pyqir/pyqir/_native.pyi b/pyqir/pyqir/_native.pyi
index 3b279e39..7cf86415 100644
--- a/pyqir/pyqir/_native.pyi
+++ b/pyqir/pyqir/_native.pyi
@@ -768,27 +768,6 @@ def const(ty: Type, value: Union[bool, int, float]) -> Constant:
     """
     ...
 
-def entry_point(
-    module: Module,
-    name: str,
-    required_num_qubits: int,
-    required_num_results: int,
-    qir_profiles: Optional[str] = "custom",
-    output_labeling_schema: Optional[str] = "",
-) -> Function:
-    """
-    Creates an entry point.
-
-    :param Module module: The parent module.
-    :param str name: The entry point name.
-    :param int required_num_qubits: The number of qubits required by the entry point.
-    :param int required_num_results: The number of results required by the entry point.
-    :param Optional[str] qir_profiles: Value identifying the profile the entry point has been compiled for. Use base_profile when QIR is compliant.
-    :param Optional[str] output_labeling_schema: An arbitrary string value that identifies the schema used by a compiler frontend that produced the IR to label the recorded output
-    :returns: An entry point.
-    """
-    ...
-
 def qir_major_version(module: Module) -> Optional[int]:
     """The QIR major version this module is built for. None if unspecified."""
     ...
@@ -1194,3 +1173,15 @@ def if_result(
     :param zero: A callable that inserts instructions for the branch where the result is zero.
     """
     ...
+
+def add_string_attribute(
+    function: Function, kind: str, value: Optional[str] = None
+) -> bool:
+    """
+    Adds a string attribute to the given function.
+
+    :param function: The function.
+    :param key: The attribute key.
+    :param value: The attribute value.
+    """
+    ...
diff --git a/pyqir/src/python.rs b/pyqir/src/python.rs
index a594d69a..88216e04 100644
--- a/pyqir/src/python.rs
+++ b/pyqir/src/python.rs
@@ -19,11 +19,11 @@ use crate::{
         PointerType, StructType, Type,
     },
     values::{
-        dynamic_qubit_management, dynamic_result_management, entry_point, extract_byte_string,
-        global_byte_string, is_entry_point, is_interop_friendly, qir_major_version,
-        qir_minor_version, qir_module, qubit, qubit_id, r#const, required_num_qubits,
-        required_num_results, result, result_id, Attribute, AttributeList, AttributeSet,
-        BasicBlock, Constant, FloatConstant, Function, IntConstant, Value,
+        add_string_attribute, dynamic_qubit_management, dynamic_result_management,
+        extract_byte_string, global_byte_string, is_entry_point, is_interop_friendly,
+        qir_major_version, qir_minor_version, qir_module, qubit, qubit_id, r#const,
+        required_num_qubits, required_num_results, result, result_id, Attribute, AttributeList,
+        AttributeSet, BasicBlock, Constant, FloatConstant, Function, IntConstant, Value,
     },
 };
 use pyo3::prelude::*;
@@ -62,9 +62,9 @@ fn _native(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<Switch>()?;
     m.add_class::<Type>()?;
     m.add_class::<Value>()?;
+    m.add_function(wrap_pyfunction!(add_string_attribute, m)?)?;
     m.add_function(wrap_pyfunction!(dynamic_qubit_management, m)?)?;
     m.add_function(wrap_pyfunction!(dynamic_result_management, m)?)?;
-    m.add_function(wrap_pyfunction!(entry_point, m)?)?;
     m.add_function(wrap_pyfunction!(extract_byte_string, m)?)?;
     m.add_function(wrap_pyfunction!(global_byte_string, m)?)?;
     m.add_function(wrap_pyfunction!(is_entry_point, m)?)?;
diff --git a/pyqir/src/values.rs b/pyqir/src/values.rs
index c1ac4689..39e6c3cd 100644
--- a/pyqir/src/values.rs
+++ b/pyqir/src/values.rs
@@ -20,7 +20,7 @@ use pyo3::{
     exceptions::{PyKeyError, PyTypeError, PyValueError},
     prelude::*,
     pyclass::CompareOp,
-    types::{PyBytes, PyLong},
+    types::{PyBytes, PyLong, PyString},
     PyRef,
 };
 use qirlib::values;
@@ -713,43 +713,6 @@ pub(crate) fn result_id(value: &Value) -> Option<u64> {
     unsafe { values::result_id(value.as_ptr()) }
 }
 
-/// Creates an entry point.
-///
-/// :param Module module: The parent module.
-/// :param str name: The entry point name.
-/// :param int required_num_qubits: The number of qubits required by the entry point.
-/// :param int required_num_results: The number of results required by the entry point.
-/// :param str qir_profiles: Value identifying the profile the entry point has been compiled for. Use base_profile when QIR is compliant.
-/// :param str output_labeling_schema: An arbitrary string value that identifies the schema used by a compiler frontend that produced the IR to label the recorded output
-/// :returns: An entry point.
-/// :rtype: Function
-#[pyfunction]
-#[pyo3(
-    text_signature = "(module, name, required_num_qubits, required_num_results, qir_profiles, output_labeling_schema)"
-)]
-pub(crate) fn entry_point(
-    py: Python,
-    module: Py<Module>,
-    name: &str,
-    required_num_qubits: u64,
-    required_num_results: u64,
-    qir_profiles: Option<&str>,
-    output_labeling_schema: Option<&str>,
-) -> PyResult<PyObject> {
-    let name = CString::new(name).unwrap();
-    unsafe {
-        let entry_point = values::entry_point(
-            module.borrow(py).as_ptr(),
-            name.as_c_str(),
-            required_num_qubits,
-            required_num_results,
-            qir_profiles.unwrap_or("custom"),
-            output_labeling_schema.unwrap_or(""),
-        );
-        Value::from_raw(py, module.into(), entry_point)
-    }
-}
-
 /// Whether the function is an entry point.
 ///
 /// :param Function function: The function.
@@ -894,3 +857,30 @@ pub(crate) fn extract_byte_string<'py>(py: Python<'py>, value: &Value) -> Option
     let string = unsafe { values::extract_string(value.as_ptr())? };
     Some(PyBytes::new(py, &string))
 }
+
+// Adds a string attribute to the given function.
+
+// :param function: The function.
+// :param kind: The attribute kind.
+// :param value: The attribute value.
+#[pyfunction]
+#[pyo3(text_signature = "(function, key, value)")]
+pub(crate) fn add_string_attribute<'py>(
+    function: PyRef<Function>,
+    key: &'py PyString,
+    value: Option<&'py PyString>,
+) {
+    let function = function.into_super().into_super().as_ptr();
+    let key = key.to_string_lossy();
+    let value = value.map(PyString::to_string_lossy);
+    unsafe {
+        values::add_string_attribute(
+            function,
+            key.as_bytes(),
+            match value {
+                Some(ref x) => x.as_bytes(),
+                None => &[],
+            },
+        );
+    }
+}
diff --git a/qirlib/src/values.rs b/qirlib/src/values.rs
index c2f516d7..70e8ac0b 100644
--- a/qirlib/src/values.rs
+++ b/qirlib/src/values.rs
@@ -38,7 +38,6 @@ pub unsafe fn result_id(value: LLVMValueRef) -> Option<u64> {
         None
     }
 }
-
 pub unsafe fn entry_point(
     module: LLVMModuleRef,
     name: &CStr,
@@ -182,12 +181,12 @@ pub unsafe fn extract_string(value: LLVMValueRef) -> Option<Vec<u8>> {
     Some(data[offset..].to_vec())
 }
 
-unsafe fn add_string_attribute(function: LLVMValueRef, kind: &[u8], value: &[u8]) {
+pub unsafe fn add_string_attribute(function: LLVMValueRef, key: &[u8], value: &[u8]) {
     let context = LLVMGetTypeContext(LLVMTypeOf(function));
     let attr = LLVMCreateStringAttribute(
         context,
-        kind.as_ptr().cast(),
-        kind.len().try_into().unwrap(),
+        key.as_ptr().cast(),
+        key.len().try_into().unwrap(),
         value.as_ptr().cast(),
         value.len().try_into().unwrap(),
     );


Test patch:
diff --git a/pyqir/tests/test_string_attributes.py b/pyqir/tests/test_string_attributes.py
new file mode 100644
index 00000000..a2f68926
--- /dev/null
+++ b/pyqir/tests/test_string_attributes.py
@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import pyqir
+from pyqir import (
+    IntType,
+    ModuleFlagBehavior,
+    Module,
+    Context,
+    add_string_attribute,
+    Function,
+    Linkage,
+    FunctionType,
+)
+import pytest
+
+
+def test_basic_key_value_attribute() -> None:
+    mod = pyqir.Module(pyqir.Context(), "test")
+    void = pyqir.Type.void(mod.context)
+    function = Function(FunctionType(void, []), Linkage.EXTERNAL, "test_function", mod)
+    add_string_attribute(function, "foo", "bar")
+    assert function.attributes.func["foo"].string_value == "bar"
+    ir = str(mod)
+    assert 'attributes #0 = { "foo"="bar" }' in ir
+
+
+def test_basic_key_only_attribute() -> None:
+    mod = pyqir.Module(pyqir.Context(), "test")
+    void = pyqir.Type.void(mod.context)
+    function = Function(FunctionType(void, []), Linkage.EXTERNAL, "test_function", mod)
+    add_string_attribute(function, "foo")
+    assert function.attributes.func["foo"].string_value == ""
+    ir = str(mod)
+    assert 'attributes #0 = { "foo" }' in ir
+
+
+def test_round_trip_serialize_parse() -> None:
+    mod = pyqir.Module(pyqir.Context(), "test")
+    void = pyqir.Type.void(mod.context)
+    function = Function(FunctionType(void, []), Linkage.EXTERNAL, "test_function", mod)
+    add_string_attribute(function, "foo", "bar")
+    # also test for non-value attributes
+    add_string_attribute(function, "entry_point", "")
+    # test behavior of empty attribute
+    add_string_attribute(function, "", "")
+    ir = str(mod)
+    parsed_mod = Module.from_ir(Context(), ir, "test")
+    assert str(parsed_mod) == str(mod)
+
+
+def test_duplicate_attr_key_replaces_previous() -> None:
+    mod = pyqir.Module(pyqir.Context(), "test")
+    void = pyqir.Type.void(mod.context)
+    function = Function(FunctionType(void, []), Linkage.EXTERNAL, "test_function", mod)
+    add_string_attribute(function, "foo", "bar")
+    add_string_attribute(function, "foo", "")
+    ir = str(mod)
+    # Tests that subsequently added attributes with the same key
+    # replace previously added ones
+    assert 'attributes #0 = { "foo" }' in ir
+
+
+def test_attribute_alphabetical_sorting() -> None:
+    mod = pyqir.Module(pyqir.Context(), "test")
+    void = pyqir.Type.void(mod.context)
+    function = Function(FunctionType(void, []), Linkage.EXTERNAL, "test_function", mod)
+    add_string_attribute(function, "b", "A")
+    add_string_attribute(function, "c", "")
+    add_string_attribute(function, "a", "a")
+    add_string_attribute(function, "1", "")
+    add_string_attribute(function, "A", "123")
+    ir = str(mod)
+    parsed_mod = Module.from_ir(Context(), ir, "test")
+    print(str(mod))
+    assert str(parsed_mod) == str(mod)
+    # Tests that attributes are sorted alphabetically by key,
+    # irrespective of their value
+    assert 'attributes #0 = { "1" "A"="123" "a"="a" "b"="A" "c" }' in ir