33# CoreML backend for delegating a EdgeProgram to CoreML.
44
55import json
6+ import logging
67
78import shutil
89import uuid
1415from typing import Any , Dict , final , List , Optional , Tuple
1516
1617import coremltools as ct
18+ import coremltools .optimize as cto
1719import executorchcoreml
1820
1921from executorch .exir .backend .backend_details import (
2325)
2426from executorch .exir .backend .compile_spec_schema import CompileSpec
2527
28+ logger = logging .getLogger (__name__ )
29+ logger .setLevel (logging .WARNING )
30+
2631
2732class COMPILE_SPEC_KEYS (Enum ):
2833 COMPUTE_UNITS = "compute_units"
2934 MODEL_TYPE = "model_type"
3035 MIN_DEPLOYMENT_TARGET = "min_deployment_target"
3136 MODEL_COMPUTE_PRECISION = "model_compute_precision"
37+ OP_LINEAR_QUANTIZER_CONFIG = "op_linear_quantizer_config"
3238
3339
3440class MODEL_PATHS (Enum ):
@@ -169,12 +175,44 @@ def generate_compute_unit_compile_spec(
169175 compute_unit .name .lower ().encode ("utf-8" ),
170176 )
171177
178+ @staticmethod
179+ def generate_op_linear_quantizer_config_compile_spec (
180+ op_linear_quantizer_config : Dict ,
181+ ) -> CompileSpec :
182+ """
183+ Returns the compile spec representing the model post conversion quantization,
184+ which is a dict that will construct cto.coreml.OpLinearQuantizerConfig
185+ """
186+ str_representation = json .dumps (op_linear_quantizer_config )
187+ byte_representation = str_representation .encode ("utf-8" )
188+ return CompileSpec (
189+ COMPILE_SPEC_KEYS .OP_LINEAR_QUANTIZER_CONFIG .value ,
190+ byte_representation ,
191+ )
192+
193+ @staticmethod
194+ def op_linear_quantizer_config_from_compile_specs (
195+ compile_specs : List [CompileSpec ],
196+ ) -> cto .coreml .OpLinearQuantizerConfig :
197+ """
198+ Returns the model's post conversion quantization by parsing the list of compile specs.
199+ """
200+ for compile_spec in compile_specs :
201+ if compile_spec .key == COMPILE_SPEC_KEYS .OP_LINEAR_QUANTIZER_CONFIG .value :
202+ config_dict_str = compile_spec .value .decode ("utf-8" )
203+ config_dict = json .loads (config_dict_str )
204+ config = cto .coreml .OpLinearQuantizerConfig ._from_dict (config_dict )
205+ return config
206+
207+ return None
208+
172209 @staticmethod
173210 def generate_compile_specs (
174211 compute_unit : ct .ComputeUnit = ct .ComputeUnit .ALL ,
175212 minimum_deployment_target : ct .target = ct .target .iOS15 ,
176213 compute_precision : ct .precision = ct .precision .FLOAT16 ,
177214 model_type : MODEL_TYPE = MODEL_TYPE .MODEL ,
215+ op_linear_quantizer_config : Optional [Dict ] = None ,
178216 ) -> List [CompileSpec ]:
179217 """
180218 Returns the list of compile specs that's used by CoreMLBackend to lower the module.
@@ -192,6 +230,12 @@ def generate_compile_specs(
192230 CoreMLBackend .generate_compute_precision_compile_spec (compute_precision )
193231 )
194232 compile_specs .append (CoreMLBackend .generate_model_type_compile_spec (model_type ))
233+ if op_linear_quantizer_config is not None :
234+ compile_specs .append (
235+ CoreMLBackend .generate_op_linear_quantizer_config_compile_spec (
236+ op_linear_quantizer_config
237+ )
238+ )
195239
196240 return compile_specs
197241
@@ -368,18 +412,18 @@ def preprocess(
368412 compile_specs ,
369413 )
370414 )
371-
372415 model_compute_precision : ct .precision = (
373416 CoreMLBackend .model_compute_precision_from_compile_specs (compile_specs )
374417 )
375-
376418 minimum_deployment_target : ct .target = (
377419 CoreMLBackend .min_deployment_target_from_compile_specs (compile_specs )
378420 )
379-
380421 compute_units : ct .ComputeUnit = CoreMLBackend .compute_unit_from_compile_specs (
381422 compile_specs
382423 )
424+ op_linear_quantizer_config = (
425+ CoreMLBackend .op_linear_quantizer_config_from_compile_specs (compile_specs )
426+ )
383427
384428 mlmodel = ct .convert (
385429 model = edge_program ,
@@ -392,4 +436,15 @@ def preprocess(
392436 compute_units = compute_units ,
393437 )
394438
439+ if op_linear_quantizer_config is not None :
440+ logger .warning (
441+ "Core ML Backend op_linear_quantizer_config API is experimental"
442+ )
443+ config = cto .coreml .OptimizationConfig (
444+ global_config = op_linear_quantizer_config ,
445+ # skip embedding
446+ op_type_configs = {"gather" : None },
447+ )
448+ mlmodel = cto .coreml .linear_quantize_weights (mlmodel , config = config )
449+
395450 return CoreMLBackend .preprocess_model (mlmodel , model_type = model_type )
0 commit comments