Allows yaml ET file to enable dtype selective build

Author: BujSet

codegen/tools/gen_oplist.py

@@ -86,59 +86,88 @@ class KernelType(IntEnum):
 
 
 def _get_operators(model_file: str) -> List[str]:
-    from executorch.codegen.tools.selective_build import (  # type: ignore[import-not-found]
-        _get_program_from_buffer,
-        _get_program_operators,
-    )
-
     print("Processing model file: ", model_file)
     with open(model_file, "rb") as f:
         buf = f.read()
+    try:
+        from executorch.codegen.tools.selective_build import (  # type: ignore[import-not-found]
+            _get_program_from_buffer,
+            _get_program_operators,
+        )
 
-    program = _get_program_from_buffer(buf)
-    operators = _get_program_operators(program)
-    print(f"Model file loaded, operators are: {operators}")
-    return operators
+        program = _get_program_from_buffer(buf)
+        operators = _get_program_operators(program)
+        print(f"Model file loaded, operators are: {operators}")
+        return operators
+    except ModuleNotFoundError:
+        from executorch.exir._serialize import _deserialize_pte_binary
+        model = _deserialize_pte_binary(buf)
+        operators = set()
+        for execPlan in model.execution_plan:
+            for op in execPlan.operators:
+                operators.add(op.name)
+        print(f"Model file loaded, operators are: {operators}")
+        return operators
 
 
 def _get_kernel_metadata_for_model(model_file: str) -> Dict[str, List[str]]:
-
-    from executorch.codegen.tools.selective_build import (  # type: ignore[import-not-found]
-        _get_io_metadata_for_program_operators,
-        _get_program_from_buffer,
-        _IOMetaData,
-    )
-
     with open(model_file, "rb") as f:
         buf = f.read()
-
-    program = _get_program_from_buffer(buf)
-    operators_with_io_metadata = _get_io_metadata_for_program_operators(program)
-
     op_kernel_key_list: Dict[str, List[str]] = {}
 
-    specialized_kernels: Set[List[_IOMetaData]]
-    for op_name, specialized_kernels in operators_with_io_metadata.items():
-        print(op_name)
-        if op_name not in op_kernel_key_list:
-            op_kernel_key_list[op_name] = []
-
-        for specialized_kernel in specialized_kernels:
-            version = "v1"
-            kernel_key = version + "/"
-            for io_metadata in specialized_kernel:
-                if io_metadata.kernel_type in [
-                    KernelType.TENSOR,
-                    KernelType.TENSOR_LIST,
-                    KernelType.OPTIONAL_TENSOR_LIST,
-                ]:
-                    dim_order = ",".join(map(str, io_metadata.dim_order))
-                    kernel_key += f"{io_metadata.dtype};{dim_order}|"
-            op_kernel_key_list[op_name].append(kernel_key[:-1])
+    try:
+        from executorch.codegen.tools.selective_build import (  # type: ignore[import-not-found]
+            _get_io_metadata_for_program_operators,
+            _get_program_from_buffer,
+            _IOMetaData,
+        )
 
+        program = _get_program_from_buffer(buf)
+        operators_with_io_metadata = _get_io_metadata_for_program_operators(program)
+
+        specialized_kernels: Set[List[_IOMetaData]]
+        for op_name, specialized_kernels in operators_with_io_metadata.items():
+            print(op_name)
+            if op_name not in op_kernel_key_list:
+                op_kernel_key_list[op_name] = []
+
+            for specialized_kernel in specialized_kernels:
+                version = "v1"
+                kernel_key = version + "/"
+                for io_metadata in specialized_kernel:
+                    if io_metadata.kernel_type in [
+                        KernelType.TENSOR,
+                        KernelType.TENSOR_LIST,
+                        KernelType.OPTIONAL_TENSOR_LIST,
+                    ]:
+                        dim_order = ",".join(map(str, io_metadata.dim_order))
+                        kernel_key += f"{io_metadata.dtype};{dim_order}|"
+                op_kernel_key_list[op_name].append(kernel_key[:-1])
+
+    except ModuleNotFoundError:
+        from executorch.exir._serialize import _deserialize_pte_binary
+        model = _deserialize_pte_binary(buf)
+        for execPlan in model.execution_plan:
+            for chain in execPlan.chains:
+                for instr in chain.instructions:
+                    op_name = execPlan.operators[instr.instr_args.op_index].name
+                    if op_name not in op_kernel_key_list:
+                        op_kernel_key_list[op_name] = []
+                    version = "v1"
+                    kernel_key = version + "/"
+                    # TODO what happens when tensors have different types withina single kernel/ is that even allowed?
+                    for tensor_arg in instr.instr_args.args:
+                        val = execPlan.values[tensor_arg].val
+
+                        # TODO is there a better way to do this?
+                        if type(val).__name__ == "Tensor":
+                            dim_order = ",".join(map(str,val.dim_order))
+                            kernel_key += f"{val.scalar_type};{dim_order}|"
+                    op_kernel_key_list[op_name].append(kernel_key[:-1])
     return op_kernel_key_list
 
 
+
 def _get_et_kernel_metadata_from_ops_yaml(ops_yaml_path: str) -> Dict[str, List[str]]:
     ops = []
     with open(ops_yaml_path, "r") as f: