Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix!: OpDef serialisation #1013

Merged
merged 3 commits into from
May 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 10 additions & 8 deletions hugr-py/src/hugr/serialization/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -500,18 +500,20 @@ class Config:
# --------------------------------------


class OpDef(BaseOp, populate_by_name=True):
class FixedHugr(ConfiguredBaseModel):
extensions: ExtensionSet
hugr: Any


class OpDef(ConfiguredBaseModel, populate_by_name=True):
"""Serializable definition for dynamically loaded operations."""

extension: ExtensionId
name: str # Unique identifier of the operation.
description: str # Human readable description of the operation.
inputs: list[tuple[str | None, Type]]
outputs: list[tuple[str | None, Type]]
misc: dict[str, Any] # Miscellaneous data associated with the operation.
def_: str | None = Field(
..., alias="def"
) # (YAML?)-encoded definition of the operation.
extension_reqs: ExtensionSet # Resources required to execute this operation.
misc: dict[str, Any] | None = None
signature: PolyFuncType | None = None
lower_funcs: list[FixedHugr]


# Now that all classes are defined, we need to update the ForwardRefs in all type
Expand Down
3 changes: 2 additions & 1 deletion hugr-py/src/hugr/serialization/testing_hugr.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from pydantic import ConfigDict
from typing import Literal
from .tys import Type, SumType, PolyFuncType, ConfiguredBaseModel, model_rebuild
from .ops import Value, OpType, classes as ops_classes
from .ops import Value, OpType, OpDef, classes as ops_classes


class TestingHugr(ConfiguredBaseModel):
Expand All @@ -14,6 +14,7 @@ class TestingHugr(ConfiguredBaseModel):
poly_func_type: PolyFuncType | None = None
value: Value | None = None
optype: OpType | None = None
op_def: OpDef | None = None

@classmethod
def get_version(cls) -> str:
Expand Down
2 changes: 1 addition & 1 deletion hugr/src/extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ pub use infer::{ExtensionSolution, InferExtensionError};

mod op_def;
pub use op_def::{
CustomSignatureFunc, CustomValidator, OpDef, SignatureFromArgs, SignatureFunc,
CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
ValidateJustArgs, ValidateTypeArgs,
};
mod type_def;
Expand Down
31 changes: 20 additions & 11 deletions hugr/src/extension/op_def.rs
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ pub struct CustomValidator {
#[serde(flatten)]
poly_func: PolyFuncType,
#[serde(skip)]
validate: Box<dyn ValidateTypeArgs>,
pub(crate) validate: Box<dyn ValidateTypeArgs>,
}

impl CustomValidator {
Expand Down Expand Up @@ -265,11 +265,19 @@ impl Debug for SignatureFunc {
/// Different ways that an [OpDef] can lower operation nodes i.e. provide a Hugr
/// that implements the operation using a set of other extensions.
#[derive(serde::Deserialize, serde::Serialize)]
#[serde(untagged)]
pub enum LowerFunc {
/// Lowering to a fixed Hugr. Since this cannot depend upon the [TypeArg]s,
/// this will generally only be applicable if the [OpDef] has no [TypeParam]s.
#[serde(rename = "hugr")]
FixedHugr(ExtensionSet, Hugr),
FixedHugr {
/// The extensions required by the [`Hugr`]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmmmm, perhaps to think about in the future, but if we insist that the Hugr is DFG-rooted, then we can extract the ExtensionSet from the root node's extension_reqs, so I'm not sure we really need this...

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, but let's leave until later.

extensions: ExtensionSet,
/// The [`Hugr`] to be used to replace [CustomOp]s matching the parent
/// [OpDef]
///
/// [CustomOp]: crate::ops::CustomOp
hugr: Hugr,
},
/// Custom binary function that can (fallibly) compute a Hugr
/// for the particular instance and set of available extensions.
#[serde(skip)]
Expand All @@ -279,7 +287,7 @@ pub enum LowerFunc {
impl Debug for LowerFunc {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::FixedHugr(_, _) => write!(f, "FixedHugr"),
Self::FixedHugr { .. } => write!(f, "FixedHugr"),
Self::CustomFunc(_) => write!(f, "<custom lower>"),
}
}
Expand All @@ -305,8 +313,7 @@ pub struct OpDef {
signature_func: SignatureFunc,
// Some operations cannot lower themselves and tools that do not understand them
// can only treat them as opaque/black-box ops.
#[serde(flatten)]
lower_funcs: Vec<LowerFunc>,
pub(crate) lower_funcs: Vec<LowerFunc>,

/// Operations can optionally implement [`ConstFold`] to implement constant folding.
#[serde(skip)]
Expand Down Expand Up @@ -360,9 +367,9 @@ impl OpDef {
self.lower_funcs
.iter()
.flat_map(|f| match f {
LowerFunc::FixedHugr(req_res, h) => {
if available_extensions.is_superset(req_res) {
Some(h.clone())
LowerFunc::FixedHugr { extensions, hugr } => {
if available_extensions.is_superset(extensions) {
Some(hugr.clone())
} else {
None
}
Expand Down Expand Up @@ -477,7 +484,6 @@ mod test {
use crate::std_extensions::collections::{EXTENSION, LIST_TYPENAME};
use crate::types::Type;
use crate::types::{type_param::TypeParam, FunctionType, PolyFuncType, TypeArg, TypeBound};
use crate::Hugr;
use crate::{const_extension_ids, Extension};

const_extension_ids! {
Expand All @@ -495,7 +501,10 @@ mod test {
let type_scheme = PolyFuncType::new(vec![TP], FunctionType::new_endo(vec![list_of_var]));

let def = e.add_op(OP_NAME, "desc".into(), type_scheme)?;
def.add_lower_func(LowerFunc::FixedHugr(ExtensionSet::new(), Hugr::default()));
def.add_lower_func(LowerFunc::FixedHugr {
extensions: ExtensionSet::new(),
hugr: crate::builder::test::simple_dfg_hugr(), // this is nonsense, but we are not testing the actual lowering here
});
def.add_misc("key", Default::default());
assert_eq!(def.description(), "desc");
assert_eq!(def.lower_funcs.len(), 1);
Expand Down
2 changes: 1 addition & 1 deletion hugr/src/hugr/serialize/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -377,7 +377,7 @@ fn roundtrip_sumtype(#[case] sum_type: SumType) {
// #[case(Value::extension(ConstF64::new(std::f64::NAN)))]
// #[case(Value::extension(ConstF64::new(std::f64::INFINITY)))]
// #[case(Value::extension(ConstF64::new(std::f64::NEG_INFINITY)))]
#[case(Value::extension(ConstF64::new(std::f64::MIN_POSITIVE)))]
#[case(Value::extension(ConstF64::new(f64::MIN_POSITIVE)))]
#[case(Value::sum(1,[Value::extension(ConstInt::new_u(2,1).unwrap())], SumType::new([vec![], vec![INT_TYPES[2].clone()]])).unwrap())]
#[case(Value::tuple([Value::false_val(), Value::extension(ConstInt::new_s(2,1).unwrap())]))]
#[case(Value::function(crate::builder::test::simple_dfg_hugr()).unwrap())]
Expand Down
88 changes: 88 additions & 0 deletions specification/schema/testing_hugr_schema_strict_v1.json
Original file line number Diff line number Diff line change
Expand Up @@ -817,6 +817,27 @@
"title": "ExtensionsParam",
"type": "object"
},
"FixedHugr": {
"additionalProperties": false,
"properties": {
"extensions": {
"items": {
"type": "string"
},
"title": "Extensions",
"type": "array"
},
"hugr": {
"title": "Hugr"
}
},
"required": [
"extensions",
"hugr"
],
"title": "FixedHugr",
"type": "object"
},
"FuncDecl": {
"additionalProperties": false,
"description": "External function declaration, linked at runtime.",
Expand Down Expand Up @@ -1363,6 +1384,62 @@
"title": "Noop",
"type": "object"
},
"OpDef": {
"additionalProperties": false,
"description": "Serializable definition for dynamically loaded operations.",
"properties": {
"extension": {
"title": "Extension",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"title": "Description",
"type": "string"
},
"misc": {
"anyOf": [
{
"type": "object"
},
{
"type": "null"
}
],
"default": null,
"title": "Misc"
},
"signature": {
"anyOf": [
{
"$ref": "#/$defs/PolyFuncType"
},
{
"type": "null"
}
],
"default": null
},
"lower_funcs": {
"items": {
"$ref": "#/$defs/FixedHugr"
},
"title": "Lower Funcs",
"type": "array"
}
},
"required": [
"extension",
"name",
"description",
"lower_funcs"
],
"title": "OpDef",
"type": "object"
},
"OpType": {
"description": "A constant operation.",
"discriminator": {
Expand Down Expand Up @@ -2325,6 +2402,17 @@
}
],
"default": null
},
"op_def": {
"anyOf": [
{
"$ref": "#/$defs/OpDef"
},
{
"type": "null"
}
],
"default": null
}
},
"title": "TestingHugr",
Expand Down
88 changes: 88 additions & 0 deletions specification/schema/testing_hugr_schema_v1.json
Original file line number Diff line number Diff line change
Expand Up @@ -817,6 +817,27 @@
"title": "ExtensionsParam",
"type": "object"
},
"FixedHugr": {
"additionalProperties": true,
"properties": {
"extensions": {
"items": {
"type": "string"
},
"title": "Extensions",
"type": "array"
},
"hugr": {
"title": "Hugr"
}
},
"required": [
"extensions",
"hugr"
],
"title": "FixedHugr",
"type": "object"
},
"FuncDecl": {
"additionalProperties": true,
"description": "External function declaration, linked at runtime.",
Expand Down Expand Up @@ -1363,6 +1384,62 @@
"title": "Noop",
"type": "object"
},
"OpDef": {
"additionalProperties": true,
"description": "Serializable definition for dynamically loaded operations.",
"properties": {
"extension": {
"title": "Extension",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"title": "Description",
"type": "string"
},
"misc": {
"anyOf": [
{
"type": "object"
},
{
"type": "null"
}
],
"default": null,
"title": "Misc"
},
"signature": {
"anyOf": [
{
"$ref": "#/$defs/PolyFuncType"
},
{
"type": "null"
}
],
"default": null
},
"lower_funcs": {
"items": {
"$ref": "#/$defs/FixedHugr"
},
"title": "Lower Funcs",
"type": "array"
}
},
"required": [
"extension",
"name",
"description",
"lower_funcs"
],
"title": "OpDef",
"type": "object"
},
"OpType": {
"description": "A constant operation.",
"discriminator": {
Expand Down Expand Up @@ -2325,6 +2402,17 @@
}
],
"default": null
},
"op_def": {
"anyOf": [
{
"$ref": "#/$defs/OpDef"
},
{
"type": "null"
}
],
"default": null
}
},
"title": "TestingHugr",
Expand Down
Loading