femtocrux 2.5.1__py3-none-any.whl → 3.0.0__py3-none-any.whl

femtocrux/VERSION

@@ -1 +1 @@
-2.5.1
+3.0.0

femtocrux/__init__.py

@@ -1,4 +1,10 @@
-from .client.client import CompilerClient, TFLiteModel, FQIRModel, ManagedCompilerClient
+from .client.client import (
+    CompilerClient,
+    TFLiteModel,
+    FQIRModel,
+    ModelAndMetadata,
+    ManagedCompilerClient,
+)
 from .version import __version__
 
 # PEP 8 definiton of public API
@@ -7,6 +13,7 @@ __all__ = [
     "CompilerClient",
     "TFLiteModel",
     "FQIRModel",
+    "ModelAndMetadata",
     "__version__",
     "ManagedCompilerClient",
 ]

femtocrux/client/client.py

@@ -15,6 +15,7 @@ from contextlib import contextmanager
 import subprocess
 
 from fmot.fqir import GraphProto
+from femtorun import IOSpec
 
 from femtocrux.util.utils import (
     read_secret_raw,
@@ -23,6 +24,8 @@ from femtocrux.util.utils import (
     deserialize_simulation_output,
 )
 
+from google.protobuf.message import Message
+
 # GRPC artifacts
 import femtocrux.grpc.compiler_service_pb2 as cs_pb2
 import femtocrux.grpc.compiler_service_pb2_grpc as cs_pb2_grpc
@@ -104,13 +107,16 @@ class Model:
     :class:`~femtocrux.client.client.TFLiteModel`.
     """
 
-    def _get_message(self, options: dict = {}) -> cs_pb2.model:
+    def _get_message(self, user_configuration: dict = {}) -> cs_pb2.model_and_metadata:
         # Format the options
-        options_struct = google.protobuf.struct_pb2.Struct()
-        options_struct.update(options)
+        user_configuration_struct = google.protobuf.struct_pb2.Struct()
+        user_configuration_struct.update(user_configuration)
 
         # Construct the model with IR
-        return cs_pb2.model(**{self._ir_name: self._get_ir()}, options=options_struct)
+        return cs_pb2.model_and_metadata(
+            **{self._ir_name: self._get_ir()},
+            user_configuration=user_configuration_struct,
+        )
 
     @property
     def _ir_name(self) -> str:
@@ -186,14 +192,54 @@ class TFLiteModel(Model):
         return cs_pb2.tflite(model=self.flatbuffer, signature_name=self.signature_name)
 
 
+@dataclass
+class ModelAndMetadata(Model):
+    """
+    Wraps an FQIR model and metadata to be compiled by femtocrux.
+
+    model: FQIRModel object
+    iospec: IOSpec object
+    opt_profile: (str) simple compiler configuration setting
+    user_configuration: (dict) fine grained compiler configuration
+    memory_reservations: (dict) requested memory to be reserved for IO buffering.
+                         Currently unsupported.
+    """
+
+    model: FQIRModel = None
+    iospec: IOSpec = None
+    opt_profile: str | None = None
+    user_configuration: dict | None = None
+    memory_reservations: dict | None = None  # Currently unsupported by the compiler
+
+    @property
+    def _ir_name(self) -> str:
+        return "fqir"
+
+    def _get_ir(self) -> Any:
+        # Send the serialized model
+        serialized_iospec = self.iospec.serialize_to_string()
+        return cs_pb2.model_and_metadata(
+            fqir=self.model._get_ir(),
+            iospec=serialized_iospec,
+            opt_profile=self.opt_profile,
+            user_configuration=self.user_configuration,
+            memory_reservations=self.memory_reservations,
+        )
+
+
 class Simulator:
     """
     Simulates a compiled model's behavior on the SPU.
     """
 
-    def __init__(self, client: "CompilerClient", model: Model, options: dict = {}):
+    def __init__(
+        self,
+        client: "CompilerClient",
+        model_and_metadata: ModelAndMetadata,
+        options: dict = {},
+    ):
         self.client = client
-        self.model = model
+        self.model_and_metadata = model_and_metadata
 
         # Create an event stream fed by a queue
         self.request_queue = queue.SimpleQueue()
@@ -201,8 +247,9 @@ class Simulator:
         self.response_iterator = client._simulate(request_iterator)
 
         # Compile the model with the first message
-        model_msg = model._get_message(options)
-        simulation_start_msg = cs_pb2.simulation_input(model=model_msg)
+        # model_msg = self.model._get_message(user_configuration=options)
+        message = model_and_metadata._get_ir()
+        simulation_start_msg = cs_pb2.simulation_input(model=message)
         self._send_request(simulation_start_msg)
 
         # Check compilation status
@@ -291,23 +338,63 @@ class CompilerClientImpl:
             server_version,
         )
 
-    def compile(self, model: Model, options: dict = {}) -> bytes:
+    def compile(
+        self,
+        model: GraphProto | Model,
+        iospec: IOSpec | None = None,
+        opt_profile: str | None = None,
+        user_configuration: dict | None = None,
+        memory_reservations: dict | None = None,
+    ) -> tuple[Message, IOSpec, Message]:
         """
         Compile the model to a bitstream.
 
-        :type model: Model, required
-        :param model: The model to be compiled.
-
-        :type options: dict, optional
-        :param options: Complier options.
-
-        :rtype: bytes
-        :return: A zip archive of compiler artifacts.
+        :model: FQIRModel or GraphProto of the model to compile
+        :iospec: IOSpec object
+        :opt_profile: (str) simple compiler configuration setting
+        :user_configuration: (dict) fine grained compiler configuration
+        :memory_reservations: (dict) requested memory to be reserved for IO buffering.
+                              Currently unsupported.
+
+        :return:
+            ProtoMessage of A zip archive of compiler artifacts.
+            ProtoMessage of iospec
+            ProtoMessage of memory_reservations
         """
 
-        response = self.stub.compile(model._get_message(options))
+        # response = self.stub.compile(model._get_message(options))
+
+        # Send the serialized model
+        if isinstance(model, GraphProto):
+            fqir_message = FQIRModel(
+                graph_proto=model,
+                batch_dim=0,
+                sequence_dim=1,
+            )
+        elif isinstance(model, FQIRModel):
+            fqir_message = model
+        else:
+            raise Exception(
+                f"Expected fqir GraphProto or protobuf Model message. "
+                f"Model allows for specifying custom batch and sequence dimensions. "
+                f"Got {type(model)}"
+            )
+        if iospec is None:
+            raise Exception(
+                "No iospec was provided. Please provide an iospec to the compiler."
+            )
+        serialized_iospec = iospec.serialize_to_string()
+        data_to_send = cs_pb2.model_and_metadata(
+            fqir=fqir_message._get_ir(),
+            iospec=serialized_iospec,
+            opt_profile=opt_profile,
+            user_configuration=user_configuration,
+            memory_reservations=memory_reservations,
+        )
+        response = self.stub.compile(data_to_send)
+        returned_iospec = IOSpec.deserialize_from_string(response.iospec)
         self._check_status(response.status)
-        return response.bitfile
+        return response.bitfile, returned_iospec, response.memory_reservations
 
     def _ping(self, message: bytes) -> None:
         """Pings the server with a message."""
@@ -319,7 +406,9 @@ class CompilerClientImpl:
         """Calls the 'simulator' bidirectional streaming RPC."""
         return self.stub.simulate(in_stream)
 
-    def simulate(self, model: Model, options: dict = {}) -> Simulator:
+    def simulate(
+        self, model_and_metadata: ModelAndMetadata, options: dict = {}
+    ) -> Simulator:
         """
         Get a simulator for the model.
 
@@ -332,10 +421,14 @@ class CompilerClientImpl:
         :rtype: Simulator
         :return: A simulator for the model.
         """
-        return Simulator(client=self, model=model, options=options)
+        return Simulator(
+            client=self, model_and_metadata=model_and_metadata, options=options
+        )
 
-    def get_simulator_object(self, model: Model, options: dict = {}) -> Simulator:
-        return self.simulate(model, options)
+    def get_simulator_object(
+        self, model_and_metadata: ModelAndMetadata, options: dict = {}
+    ) -> Simulator:
+        return self.simulate(model_and_metadata, options)
 
     def _server_version(self) -> str:
         """Queries the femtocrux version running on the server."""

femtocrux/grpc/compiler_service_pb2.py

@@ -26,7 +26,7 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16\x63ompiler_service.proto\x12\nfscompiler\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1cgoogle/protobuf/struct.proto\"g\n\x04\x66qir\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x16\n\tbatch_dim\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12\x19\n\x0csequence_dim\x18\x03 \x01(\x03H\x01\x88\x01\x01\x42\x0c\n\n_batch_dimB\x0f\n\r_sequence_dim\"G\n\x06tflite\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x1b\n\x0esignature_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\x11\n\x0f_signature_name\"\x90\x01\n\x05model\x12 \n\x04\x66qir\x18\x01 \x01(\x0b\x32\x10.fscompiler.fqirH\x00\x12$\n\x06tflite\x18\x02 \x01(\x0b\x32\x12.fscompiler.tfliteH\x00\x12-\n\x07options\x18\x03 \x01(\x0b\x32\x17.google.protobuf.StructH\x01\x88\x01\x01\x42\x04\n\x02irB\n\n\x08_options\"&\n\x06status\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0b\n\x03msg\x18\x02 \x01(\t\"I\n\x12\x63ompiled_artifacts\x12\x0f\n\x07\x62itfile\x18\x01 \x01(\x0c\x12\"\n\x06status\x18\x02 \x01(\x0b\x32\x12.fscompiler.status\"5\n\x07ndarray\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x03\x12\r\n\x05\x64type\x18\x03 \x01(\t\"\x14\n\x04\x64\x61ta\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\xe6\x01\n\x0fsimulation_data\x12\x37\n\x06inputs\x18\x01 \x03(\x0b\x32\'.fscompiler.simulation_data.InputsEntry\x12\x19\n\x0csim_duration\x18\x02 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cinput_period\x18\x03 \x01(\x02H\x01\x88\x01\x01\x1a\x42\n\x0bInputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.fscompiler.ndarray:\x02\x38\x01\x42\x0f\n\r_sim_durationB\x0f\n\r_input_period\"t\n\x10simulation_input\x12\"\n\x05model\x18\x01 \x01(\x0b\x32\x11.fscompiler.modelH\x00\x12+\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x1b.fscompiler.simulation_dataH\x00\x42\x0f\n\rmodel_or_data\"\xc9\x01\n\x11simulation_output\x12;\n\x07outputs\x18\x01 \x03(\x0b\x32*.fscompiler.simulation_output.OutputsEntry\x12\x0e\n\x06report\x18\x02 \x01(\t\x12\"\n\x06status\x18\x03 \x01(\x0b\x32\x12.fscompiler.status\x1a\x43\n\x0cOutputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.fscompiler.ndarray:\x02\x38\x01\"\x1f\n\x0cversion_info\x12\x0f\n\x07version\x18\x01 \x01(\t2\x85\x02\n\x07\x43ompile\x12>\n\x07\x63ompile\x12\x11.fscompiler.model\x1a\x1e.fscompiler.compiled_artifacts\"\x00\x12,\n\x04ping\x12\x10.fscompiler.data\x1a\x10.fscompiler.data\"\x00\x12M\n\x08simulate\x12\x1c.fscompiler.simulation_input\x1a\x1d.fscompiler.simulation_output\"\x00(\x01\x30\x01\x12=\n\x07version\x12\x16.google.protobuf.Empty\x1a\x18.fscompiler.version_info\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16\x63ompiler_service.proto\x12\nfscompiler\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1cgoogle/protobuf/struct.proto\"g\n\x04\x66qir\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x16\n\tbatch_dim\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12\x19\n\x0csequence_dim\x18\x03 \x01(\x03H\x01\x88\x01\x01\x42\x0c\n\n_batch_dimB\x0f\n\r_sequence_dim\"G\n\x06tflite\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x1b\n\x0esignature_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\x11\n\x0f_signature_name\"\xd0\x02\n\x12model_and_metadata\x12 \n\x04\x66qir\x18\x01 \x01(\x0b\x32\x10.fscompiler.fqirH\x00\x12$\n\x06tflite\x18\x02 \x01(\x0b\x32\x12.fscompiler.tfliteH\x00\x12\x13\n\x06iospec\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x18\n\x0bopt_profile\x18\x04 \x01(\tH\x02\x88\x01\x01\x12\x38\n\x12user_configuration\x18\x05 \x01(\x0b\x32\x17.google.protobuf.StructH\x03\x88\x01\x01\x12\x39\n\x13memory_reservations\x18\x06 \x01(\x0b\x32\x17.google.protobuf.StructH\x04\x88\x01\x01\x42\x04\n\x02irB\t\n\x07_iospecB\x0e\n\x0c_opt_profileB\x15\n\x13_user_configurationB\x16\n\x14_memory_reservations\"&\n\x06status\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0b\n\x03msg\x18\x02 \x01(\t\"\xbc\x01\n\x12\x63ompiled_artifacts\x12\x0f\n\x07\x62itfile\x18\x01 \x01(\x0c\x12\x13\n\x06iospec\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x39\n\x13memory_reservations\x18\x03 \x01(\x0b\x32\x17.google.protobuf.StructH\x01\x88\x01\x01\x12\"\n\x06status\x18\x04 \x01(\x0b\x32\x12.fscompiler.statusB\t\n\x07_iospecB\x16\n\x14_memory_reservations\"5\n\x07ndarray\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x03\x12\r\n\x05\x64type\x18\x03 \x01(\t\"\x14\n\x04\x64\x61ta\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\xe6\x01\n\x0fsimulation_data\x12\x37\n\x06inputs\x18\x01 \x03(\x0b\x32\'.fscompiler.simulation_data.InputsEntry\x12\x19\n\x0csim_duration\x18\x02 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cinput_period\x18\x03 \x01(\x02H\x01\x88\x01\x01\x1a\x42\n\x0bInputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.fscompiler.ndarray:\x02\x38\x01\x42\x0f\n\r_sim_durationB\x0f\n\r_input_period\"\x81\x01\n\x10simulation_input\x12/\n\x05model\x18\x01 \x01(\x0b\x32\x1e.fscompiler.model_and_metadataH\x00\x12+\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x1b.fscompiler.simulation_dataH\x00\x42\x0f\n\rmodel_or_data\"\xc9\x01\n\x11simulation_output\x12;\n\x07outputs\x18\x01 \x03(\x0b\x32*.fscompiler.simulation_output.OutputsEntry\x12\x0e\n\x06report\x18\x02 \x01(\t\x12\"\n\x06status\x18\x03 \x01(\x0b\x32\x12.fscompiler.status\x1a\x43\n\x0cOutputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.fscompiler.ndarray:\x02\x38\x01\"\x1f\n\x0cversion_info\x12\x0f\n\x07version\x18\x01 \x01(\t2\x92\x02\n\x07\x43ompile\x12K\n\x07\x63ompile\x12\x1e.fscompiler.model_and_metadata\x1a\x1e.fscompiler.compiled_artifacts\"\x00\x12,\n\x04ping\x12\x10.fscompiler.data\x1a\x10.fscompiler.data\"\x00\x12M\n\x08simulate\x12\x1c.fscompiler.simulation_input\x1a\x1d.fscompiler.simulation_output\"\x00(\x01\x30\x01\x12=\n\x07version\x12\x16.google.protobuf.Empty\x1a\x18.fscompiler.version_info\"\x00\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -41,28 +41,28 @@ if not _descriptor._USE_C_DESCRIPTORS:
   _globals['_FQIR']._serialized_end=200
   _globals['_TFLITE']._serialized_start=202
   _globals['_TFLITE']._serialized_end=273
-  _globals['_MODEL']._serialized_start=276
-  _globals['_MODEL']._serialized_end=420
-  _globals['_STATUS']._serialized_start=422
-  _globals['_STATUS']._serialized_end=460
-  _globals['_COMPILED_ARTIFACTS']._serialized_start=462
-  _globals['_COMPILED_ARTIFACTS']._serialized_end=535
-  _globals['_NDARRAY']._serialized_start=537
-  _globals['_NDARRAY']._serialized_end=590
-  _globals['_DATA']._serialized_start=592
-  _globals['_DATA']._serialized_end=612
-  _globals['_SIMULATION_DATA']._serialized_start=615
-  _globals['_SIMULATION_DATA']._serialized_end=845
-  _globals['_SIMULATION_DATA_INPUTSENTRY']._serialized_start=745
-  _globals['_SIMULATION_DATA_INPUTSENTRY']._serialized_end=811
-  _globals['_SIMULATION_INPUT']._serialized_start=847
-  _globals['_SIMULATION_INPUT']._serialized_end=963
-  _globals['_SIMULATION_OUTPUT']._serialized_start=966
-  _globals['_SIMULATION_OUTPUT']._serialized_end=1167
-  _globals['_SIMULATION_OUTPUT_OUTPUTSENTRY']._serialized_start=1100
-  _globals['_SIMULATION_OUTPUT_OUTPUTSENTRY']._serialized_end=1167
-  _globals['_VERSION_INFO']._serialized_start=1169
-  _globals['_VERSION_INFO']._serialized_end=1200
-  _globals['_COMPILE']._serialized_start=1203
-  _globals['_COMPILE']._serialized_end=1464
+  _globals['_MODEL_AND_METADATA']._serialized_start=276
+  _globals['_MODEL_AND_METADATA']._serialized_end=612
+  _globals['_STATUS']._serialized_start=614
+  _globals['_STATUS']._serialized_end=652
+  _globals['_COMPILED_ARTIFACTS']._serialized_start=655
+  _globals['_COMPILED_ARTIFACTS']._serialized_end=843
+  _globals['_NDARRAY']._serialized_start=845
+  _globals['_NDARRAY']._serialized_end=898
+  _globals['_DATA']._serialized_start=900
+  _globals['_DATA']._serialized_end=920
+  _globals['_SIMULATION_DATA']._serialized_start=923
+  _globals['_SIMULATION_DATA']._serialized_end=1153
+  _globals['_SIMULATION_DATA_INPUTSENTRY']._serialized_start=1053
+  _globals['_SIMULATION_DATA_INPUTSENTRY']._serialized_end=1119
+  _globals['_SIMULATION_INPUT']._serialized_start=1156
+  _globals['_SIMULATION_INPUT']._serialized_end=1285
+  _globals['_SIMULATION_OUTPUT']._serialized_start=1288
+  _globals['_SIMULATION_OUTPUT']._serialized_end=1489
+  _globals['_SIMULATION_OUTPUT_OUTPUTSENTRY']._serialized_start=1422
+  _globals['_SIMULATION_OUTPUT_OUTPUTSENTRY']._serialized_end=1489
+  _globals['_VERSION_INFO']._serialized_start=1491
+  _globals['_VERSION_INFO']._serialized_end=1522
+  _globals['_COMPILE']._serialized_start=1525
+  _globals['_COMPILE']._serialized_end=1799
 # @@protoc_insertion_point(module_scope)

femtocrux/grpc/compiler_service_pb2_grpc.py

@@ -38,7 +38,7 @@ class CompileStub(object):
         """
         self.compile = channel.unary_unary(
                 '/fscompiler.Compile/compile',
-                request_serializer=compiler__service__pb2.model.SerializeToString,
+                request_serializer=compiler__service__pb2.model_and_metadata.SerializeToString,
                 response_deserializer=compiler__service__pb2.compiled_artifacts.FromString,
                 _registered_method=True)
         self.ping = channel.unary_unary(
@@ -91,7 +91,7 @@ def add_CompileServicer_to_server(servicer, server):
     rpc_method_handlers = {
             'compile': grpc.unary_unary_rpc_method_handler(
                     servicer.compile,
-                    request_deserializer=compiler__service__pb2.model.FromString,
+                    request_deserializer=compiler__service__pb2.model_and_metadata.FromString,
                     response_serializer=compiler__service__pb2.compiled_artifacts.SerializeToString,
             ),
             'ping': grpc.unary_unary_rpc_method_handler(
@@ -136,7 +136,7 @@ class Compile(object):
             request,
             target,
             '/fscompiler.Compile/compile',
-            compiler__service__pb2.model.SerializeToString,
+            compiler__service__pb2.model_and_metadata.SerializeToString,
             compiler__service__pb2.compiled_artifacts.FromString,
             options,
             channel_credentials,

femtocrux/server/compiler_frontend.py

@@ -13,10 +13,11 @@ import zipfile
 
 from fmot import ConvertedModel
 from fmot.fqir import GraphProto
-from femtomapper import MapperConf, Mapper, MapperState
+from femtomapper import compile as femtomapper_compile
 from femtobehav.fasmir import FASMIR
 from femtobehav.sim import SimRunner
-from typing import Any
+from femtorun import IOSpec
+from typing import Any, Literal
 
 
 class CompilerFrontend:
@@ -25,13 +26,15 @@ class CompilerFrontend:
     def __init__(self, input_ir: Any, fasmir: FASMIR = None):
         self.input_ir = input_ir
         self.fasmir = fasmir
+        self.iospec = None
+        self.memory_reservations = None
         # self.io_wrapper = io_wrapper
 
     @property
     def is_compiled(self):
         return self.fasmir is not None
 
-    def _compile(self, input_ir: Any, options: dict) -> FASMIR:
+    def _compile(self, input_ir: Any) -> FASMIR:
         """
         Runs FM compiler to generate FASMIR, and encode io information in a
         SimIOWrapper object.
@@ -44,9 +47,19 @@ class CompilerFrontend:
             "Subclasses need to implement this based on their input ir"
         )
 
-    def compile(self, options: dict = {}):
+    def compile(
+        self,
+        input_ir: GraphProto,
+        iospec: IOSpec | None = None,
+        opt_profile: Literal["efficient", "fast", "conservative"] | None = "efficient",
+        user_configuration: dict | None = None,
+        memory_reservations: dict | None = None,
+    ) -> tuple[FASMIR, IOSpec, dict]:
         if not self.is_compiled:
-            self.fasmir = self._compile(self.input_ir, options)
+            self.fasmir, self.iospec, self.memory_reservations = self._compile(
+                input_ir, iospec, opt_profile, user_configuration, memory_reservations
+            )
+        return self.fasmir, self.iospec, self.memory_reservations
 
     def dump_bitfile(self, encrypt: bool = True) -> bytes:
         """Dumps a bitfile used to program the SPU."""
@@ -56,7 +69,9 @@ class CompilerFrontend:
         with tempfile.TemporaryFile() as tmpfile:
             with tempfile.TemporaryDirectory() as dirname:
                 # Dump memory files to a directory
-                runner = SimRunner(self.fasmir, data_dir=dirname, encrypt=encrypt)
+                runner = SimRunner(
+                    self.fasmir, io_spec=self.iospec, data_dir=dirname, encrypt=encrypt
+                )
                 runner.reset()
                 runner.finish()
 
@@ -84,7 +99,11 @@ class CompilerFrontend:
         return fasmir_var.numpy.shape[0]
 
     def run_behavioral_simulator(
-        self, inputs: dict[str, np.ndarray], input_period: float = None, **kwargs
+        self,
+        inputs: dict[str, np.ndarray],
+        input_period: float = None,
+        iospec: IOSpec | None = None,
+        **kwargs
     ):
         """
         Runs the behavioral simulator and returns outputs and metrics.
@@ -96,7 +115,7 @@ class CompilerFrontend:
             input_period (float, optional): total simulation time.
 
         """
-        runner = SimRunner(self.fasmir, **kwargs)
+        runner = SimRunner(self.fasmir, io_spec=iospec, **kwargs)
         runner.reset()
         outputs, __, __ = runner.run(inputs)
         metrics = runner.get_metrics(input_period, concise=True, as_yamlable=True)
@@ -104,19 +123,6 @@ class CompilerFrontend:
         return outputs, metrics
 
 
-def _compile_fqir(graph: GraphProto, options: dict) -> FASMIR:
-    mapper_conf = MapperConf(**options)
-    mapper = Mapper(mapper_conf)
-    mapper_state = MapperState(fqir=graph)
-
-    # compile:
-    mapper_state = mapper.do(mapper_state)
-
-    # extract fasmir
-    fasmir = mapper_state.fasmir
-    return fasmir
-
-
 class TorchCompiler(CompilerFrontend):
     def __init__(self, graph: GraphProto, batch_dim: int = None, seq_dim: int = None):
         assert isinstance(graph, GraphProto)
@@ -125,10 +131,19 @@ class TorchCompiler(CompilerFrontend):
         self.batch_dim = batch_dim
         self.seq_dim = seq_dim
 
-    def _compile(self, input_ir: GraphProto, options: dict) -> FASMIR:
-        fasmir = _compile_fqir(input_ir, options)
-        # wrapper = self._get_fqir_iowrapper(input_ir, fasmir)
-        return fasmir
+    def _compile(
+        self,
+        input_ir: GraphProto,
+        iospec: IOSpec | None = None,
+        opt_profile: Literal["efficient", "fast", "conservative"] | None = "efficient",
+        user_configuration: dict | None = None,
+        memory_reservations: dict | None = None,
+    ) -> tuple[FASMIR, IOSpec, dict]:
+        mapper_state, iospec, memory_reservations = femtomapper_compile(
+            input_ir, iospec, opt_profile, user_configuration, memory_reservations
+        )
+        fasmir = mapper_state.fasmir
+        return fasmir, iospec, memory_reservations
 
     @classmethod
     def from_fqir(cls, graph: GraphProto, batch_dim: int = None, seq_dim: int = None):

femtocrux/server/server.py

@@ -1,12 +1,12 @@
 import argparse
 from collections.abc import Iterable
 import concurrent
-from google.protobuf.json_format import MessageToDict
 import grpc
 import logging
 import pickle
 import sys
 
+from femtorun import IOSpec
 from femtocrux.server import CompilerFrontend, TorchCompiler
 
 from femtocrux.server.exceptions import format_exception, format_exception_from_exc
@@ -36,8 +36,9 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
         self.logger = logging.getLogger("CompileServicer")
         self.logger.setLevel(logging.DEBUG)
         self.logger.info("Starting compile server.")
+        self.iospec = None
 
-    def _get_fqir_compiler(self, model: cs_pb2.model) -> CompilerFrontend:
+    def _get_fqir_compiler(self, model: cs_pb2.model_and_metadata) -> CompilerFrontend:
         """Get a Torch compiler from an FQIR a model message"""
         # Deserialize FQIR
         fqir = model.fqir
@@ -50,24 +51,36 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
             seq_dim=field_or_none(fqir, "sequence_dim"),
         )
 
-    def _compile_model(self, model: cs_pb2.model, context) -> CompilerFrontend:
+    def _compile_model(
+        self, model_and_metadata: cs_pb2.model_and_metadata, context
+    ) -> CompilerFrontend:
         """Compile a model, for simulation or bitfile generation."""
-        # Get a compiler for the model
-        model_type_map = {
-            "fqir": self._get_fqir_compiler,
-        }
-        model_type = model.WhichOneof("ir")
-        compiler = model_type_map[model_type](model)
-
-        # Get the compiler options
-        options_struct = field_or_none(model, "options")
-        if options_struct is None:
-            options = {}
-        else:
-            options = MessageToDict(options_struct)
+
+        # Deserialize FQIR
+        fqir = model_and_metadata.fqir
+        graph_proto = pickle.loads(fqir.model)
+
+        # Compile the FQIR
+        compiler = TorchCompiler(
+            graph_proto,
+            batch_dim=field_or_none(fqir, "batch_dim"),
+            seq_dim=field_or_none(fqir, "sequence_dim"),
+        )
+
+        # Get the iospec
+        iospec = field_or_none(model_and_metadata, "iospec")
+        iospec = IOSpec.deserialize_from_string(iospec)
+        self.iospec = iospec
+        # Get the opt_profile
+        opt_profile = field_or_none(model_and_metadata, "opt_profile")
+        # Get the user_configuration
+        user_configuration = field_or_none(model_and_metadata, "user_configuration")
+        memory_reservations = field_or_none(model_and_metadata, "memory_reservations")
 
         # Compile the model
-        compiler.compile(options=options)
+        compiler.compile(
+            graph_proto, iospec, opt_profile, user_configuration, memory_reservations
+        )
         assert compiler.is_compiled, "Expected compilation completed"
         return compiler
 
@@ -88,10 +101,15 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
             )
 
         # Return the bitfile
-        return cs_pb2.compiled_artifacts(
-            bitfile=bitfile, status=cs_pb2.status(success=True)
+        ret_message = cs_pb2.compiled_artifacts(
+            bitfile=bitfile,
+            iospec=compiler.iospec.serialize_to_string(),
+            memory_reservations=compiler.memory_reservations,
+            status=cs_pb2.status(success=True),
         )
 
+        return ret_message
+
     def ping(self, data: cs_pb2.data, context) -> cs_pb2.data:
         """Round-trip a message."""
         return data
@@ -107,7 +125,11 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
             # Check that this is a model request
             if not model_request.WhichOneof("model_or_data") == "model":
                 yield cs_pb2.simulation_output(
-                    status=cs_pb2.status(success=False, msg="Expected model message.")
+                    status=cs_pb2.status(
+                        success=False,
+                        msg=f"Expected model message, got "
+                        f"{model_request.WhichOneof('model_or_data')}",
+                    )
                 )
                 continue
 
@@ -147,6 +169,7 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
                     outputs, metrics = compiler.run_behavioral_simulator(
                         deserialized_inputs,
                         input_period=field_or_none(data, "input_period"),
+                        iospec=self.iospec,
                     )
                 else:
                     raise (Exception("Didn't get a model or data"))

{femtocrux-2.5.1.dist-info → femtocrux-3.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: femtocrux
-Version: 2.5.1
+Version: 3.0.0
 Summary: Femtosense Compiler
 Home-page: https://github.com/femtosense/femtocrux
 Author: Femtosense

femtocrux-3.0.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+femtocrux/ENV_REQUIREMENTS.sh,sha256=t_O1B4hJAMgxvH9gwp1qls6eVFmhSYBJe64KmuK_H-4,1389
+femtocrux/PY_REQUIREMENTS,sha256=UwXV0o3gieruZUYdN9r2bqQ0Wcf_vosjeP6LJVx6oF0,275
+femtocrux/VERSION,sha256=KYW-iyjTrehY6Nj7S8IvVlsb9gIN_5gtzhQfdyG5mZw,6
+femtocrux/__init__.py,sha256=VeuEHuCvQmqHgzj_ogjIbwGE_E8UFLYSJN6GTpXBfe0,409
+femtocrux/version.py,sha256=uNg2kHxQo6oUN1ah7s9_85rCZVRoTHGPD1GAQPZW4lw,164
+femtocrux/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+femtocrux/client/client.py,sha256=UrO8OYhw6KTbI_tHz_DpRwBx93MwO1MvrkMUnAQaR94,29797
+femtocrux/grpc/__init__.py,sha256=uiMHQt5I2eAKJqI3Zh0h1Gm7cmPR4PbaGS71nCJQCGw,169
+femtocrux/grpc/compiler_service_pb2.py,sha256=ChHlC3TvWNFR-umRg2l7j3FAiybe1EGBEJfYB0Z78zM,6084
+femtocrux/grpc/compiler_service_pb2_grpc.py,sha256=b_fyld4rpKE4n5MimKOoyNZ5VU4UqUS_R3EM3_daqxM,8539
+femtocrux/server/__init__.py,sha256=kpm2AGLxo3AA9zo-G8N0Dm0gmWbvAd6jfhMbtkeHcxo,77
+femtocrux/server/compiler_frontend.py,sha256=kXVg_JIQyLY9zdwT1s5xusWaBd-9t4gDwXrU6FX6UqI,6217
+femtocrux/server/exceptions.py,sha256=lI6n471n5QKf5G3aL_1kuBVEItD-jBgithVVpPDwNYc,609
+femtocrux/server/healthcheck.py,sha256=ehqAwnv0D0zpy-AUZAPwv8rp874DZCwUmP8nzdXzZvI,1565
+femtocrux/server/server.py,sha256=k7GiLHxtzo5UH1O0edav4yEPgYD682DmB1aTZodtBuQ,8268
+femtocrux/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+femtocrux/util/utils.py,sha256=v8v09aDaOtfejpSIB-LgaA--JJj9o2_BWvNwjd8wfww,4134
+femtocrux-3.0.0.dist-info/licenses/LICENSE,sha256=eN9ZI1xHjUmFvN3TEeop5kBGXRUBfbsl55KBNBYYFqI,36
+femtocrux-3.0.0.dist-info/METADATA,sha256=TZkC3EUSy5ZqBYmeGccdq4HAYx3NL_wQoYWipcKccDk,2763
+femtocrux-3.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+femtocrux-3.0.0.dist-info/top_level.txt,sha256=BkTttlioC3je__8577wxRieZqY3Abu7FOOdMnmYbcNI,10
+femtocrux-3.0.0.dist-info/RECORD,,

femtocrux-2.5.1.dist-info/RECORD

@@ -1,22 +0,0 @@
-femtocrux/ENV_REQUIREMENTS.sh,sha256=t_O1B4hJAMgxvH9gwp1qls6eVFmhSYBJe64KmuK_H-4,1389
-femtocrux/PY_REQUIREMENTS,sha256=UwXV0o3gieruZUYdN9r2bqQ0Wcf_vosjeP6LJVx6oF0,275
-femtocrux/VERSION,sha256=mWAPanv-bDPr0aJRj0SoYaZ6_EDCXaQrxiJZVxZSlYQ,6
-femtocrux/__init__.py,sha256=yIWd9I2PEXCn_PKIILAN3mkWeTf0tgtVualeTIHNxfQ,342
-femtocrux/version.py,sha256=uNg2kHxQo6oUN1ah7s9_85rCZVRoTHGPD1GAQPZW4lw,164
-femtocrux/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-femtocrux/client/client.py,sha256=eJH-CM_Pob1V-M9EuwhFyyWd1NsdLBIUiF99Psz46-8,26428
-femtocrux/grpc/__init__.py,sha256=uiMHQt5I2eAKJqI3Zh0h1Gm7cmPR4PbaGS71nCJQCGw,169
-femtocrux/grpc/compiler_service_pb2.py,sha256=grQNjlnoUg8KKcykK7JcT7Pmv_7R7sjLpsoG_EHGyoY,5494
-femtocrux/grpc/compiler_service_pb2_grpc.py,sha256=Hl1bGJEYA7CvdIj38FUvlEolPGCoWyK8WazqOhx4v8s,8500
-femtocrux/server/__init__.py,sha256=kpm2AGLxo3AA9zo-G8N0Dm0gmWbvAd6jfhMbtkeHcxo,77
-femtocrux/server/compiler_frontend.py,sha256=-pfhOHDkXWReoB273Bc9s2s7BVTMZoCSyGM7pudmMhY,5472
-femtocrux/server/exceptions.py,sha256=lI6n471n5QKf5G3aL_1kuBVEItD-jBgithVVpPDwNYc,609
-femtocrux/server/healthcheck.py,sha256=ehqAwnv0D0zpy-AUZAPwv8rp874DZCwUmP8nzdXzZvI,1565
-femtocrux/server/server.py,sha256=sumnTj63uOGkzwj4-xwTuN4wVB7CDopy8e7PLOpW64Q,7405
-femtocrux/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-femtocrux/util/utils.py,sha256=v8v09aDaOtfejpSIB-LgaA--JJj9o2_BWvNwjd8wfww,4134
-femtocrux-2.5.1.dist-info/licenses/LICENSE,sha256=eN9ZI1xHjUmFvN3TEeop5kBGXRUBfbsl55KBNBYYFqI,36
-femtocrux-2.5.1.dist-info/METADATA,sha256=8hmP0A4KIm8CLNlJA0X4GuW2MjuthJYg2NOpEg2ruZA,2763
-femtocrux-2.5.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-femtocrux-2.5.1.dist-info/top_level.txt,sha256=BkTttlioC3je__8577wxRieZqY3Abu7FOOdMnmYbcNI,10
-femtocrux-2.5.1.dist-info/RECORD,,