diff --git a/codegen/tools/gen_oplist.py b/codegen/tools/gen_oplist.py index e4dfc7c4045..dc591defda8 100644 --- a/codegen/tools/gen_oplist.py +++ b/codegen/tools/gen_oplist.py @@ -189,6 +189,23 @@ def _dump_yaml( ) +def create_kernel_key(maybe_kernel_key: str) -> str: + # It is a kernel key. + if maybe_kernel_key.lstrip().startswith("v1"): + return maybe_kernel_key + # It is a dtype. + else: + # Generate a kernel key based on the dtype provided. + # Note: no dim order is included in this kernel key. + # For a description of the kernel key format, see + # executorch/blob/main/runtime/kernel/operator_registry.h#L97-L123 + try: + dtype = ScalarType[maybe_kernel_key] + return "v1/" + str(dtype.value) + ";" + except KeyError: + raise Exception(f"Unknown dtype: {maybe_kernel_key}") + + def gen_oplist( output_path: str, model_file_path: Optional[str] = None, @@ -223,7 +240,11 @@ def gen_oplist( ops_and_metadata = json.loads(ops_dict) for op, metadata in ops_and_metadata.items(): op_set.update({op}) - op_metadata = metadata if len(metadata) > 0 else ["default"] + op_metadata = ( + [create_kernel_key(x) for x in metadata] + if len(metadata) > 0 + else ["default"] + ) et_kernel_metadata = merge_et_kernel_metadata( et_kernel_metadata, {op: op_metadata} ) diff --git a/codegen/tools/test/test_gen_oplist.py b/codegen/tools/test/test_gen_oplist.py index c47ac252f9c..f5c6829d6a0 100644 --- a/codegen/tools/test/test_gen_oplist.py +++ b/codegen/tools/test/test_gen_oplist.py @@ -13,6 +13,7 @@ import executorch.codegen.tools.gen_oplist as gen_oplist import yaml +from executorch.codegen.tools.gen_oplist import ScalarType class TestGenOpList(unittest.TestCase): @@ -89,7 +90,7 @@ def test_gen_op_list_with_root_ops_and_dtypes( ) -> None: output_path = os.path.join(self.temp_dir.name, "output.yaml") ops_dict = { - "aten::add": ["v1/3;0,1|3;0,1|3;0,1|3;0,1", "v1/6;0,1|6;0,1|6;0,1|6;0,1"], + "aten::add": ["v1/3;0,1|3;0,1|3;0,1|3;0,1", ScalarType.Float.name], "aten::mul": [], } args = [ @@ -104,7 +105,7 @@ def test_gen_op_list_with_root_ops_and_dtypes( { "aten::add": [ "v1/3;0,1|3;0,1|3;0,1|3;0,1", - "v1/6;0,1|6;0,1|6;0,1|6;0,1", + "v1/6;", ], "aten::mul": ["default"], }, diff --git a/codegen/tools/test/test_gen_selected_op_variants.py b/codegen/tools/test/test_gen_selected_op_variants.py index e6f056e1302..bdde7ac2236 100644 --- a/codegen/tools/test/test_gen_selected_op_variants.py +++ b/codegen/tools/test/test_gen_selected_op_variants.py @@ -12,7 +12,7 @@ import expecttest -class TestGenSelectedMobileOpsHeader(expecttest.TestCase): +class TestGenSelectedOpVariants(expecttest.TestCase): def setUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.addCleanup(self.temp_dir.cleanup) @@ -84,7 +84,79 @@ def test_generates_correct_header(self) -> None: ) -class TestGenSelectedMobileOpsHeader_Empty(expecttest.TestCase): +class TestGenSelectedOpVariants_UsingDtypeString(expecttest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.addCleanup(self.temp_dir.cleanup) + self.selected_ops_yaml = os.path.join( + self.temp_dir.name, "selected_operators.yaml" + ) + with open(self.selected_ops_yaml, "w") as f: + f.write( + """ +include_all_non_op_selectives: False +include_all_operators: False +debug_info: + - model1@v100 + - model2@v50 +operators: + aten::add: + is_root_operator: Yes + is_used_for_training: Yes + include_all_overloads: No + aten::add.int: + is_root_operator: No + is_used_for_training: No + include_all_overloads: Yes +kernel_metadata: {} +et_kernel_metadata: + aten::add.out: + # A list of different kernel keys (tensors with dtype-enum/dim-order) combinations used in model + - v1/6; # Float + - v1/3; # Int + aten::mul.out: + - v1/6; # Float + aten::sub.out: + - default +build_features: [] +custom_classes: [] + """ + ) + + def tearDown(self): + self.temp_dir.cleanup() + + def test_generates_correct_header(self) -> None: + gen_selected_op_variants.write_selected_op_variants( + os.path.join(self.temp_dir.name, "selected_operators.yaml"), + self.temp_dir.name, + ) + with open( + os.path.join(self.temp_dir.name, "selected_op_variants.h"), "r" + ) as result: + self.assertExpectedInline( + result.read(), + """#pragma once +/** + * Generated by executorch/codegen/tools/gen_selected_op_variants.py + */ + +inline constexpr bool should_include_kernel_dtype( + const char *operator_name, + executorch::aten::ScalarType scalar_type +) { + return ((executorch::aten::string_view(operator_name).compare("add.out") == 0) + && (scalar_type == executorch::aten::ScalarType::Float || scalar_type == executorch::aten::ScalarType::Int)) + || ((executorch::aten::string_view(operator_name).compare("mul.out") == 0) + && (scalar_type == executorch::aten::ScalarType::Float)) + || ((executorch::aten::string_view(operator_name).compare("sub.out") == 0) + && (true)); +} +""", + ) + + +class TestGenSelectedOpVariants_Empty(expecttest.TestCase): def setUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.addCleanup(self.temp_dir.cleanup) diff --git a/examples/selective_build/targets.bzl b/examples/selective_build/targets.bzl index 173ab63fe44..276ee3afe41 100644 --- a/examples/selective_build/targets.bzl +++ b/examples/selective_build/targets.bzl @@ -1,5 +1,5 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_oss_build_kwargs", "is_xplat", "runtime") -load("@fbsource//xplat/executorch/codegen:codegen.bzl", "et_operator_library", "executorch_generated_lib") +load("@fbsource//xplat/executorch/codegen:codegen.bzl", "et_operator_library", "executorch_generated_lib", "ScalarType") def define_common_targets(): """Defines targets that should be shared between fbcode and xplat. @@ -49,7 +49,9 @@ def define_common_targets(): et_operator_library( name = "select_ops_in_dict", ops_dict = { - "aten::add.out": ["v1/3;0,1", "v1/6;0,1"], # int, float + # 1. Use kernel key, generated with a model, or + # 2. Specify the dtype, from executorch/codegen/codegen.bzl + "aten::add.out": ["v1/3;0,1", ScalarType("Float")], # int, float "aten::mm.out": [], # all dtypes }, ) diff --git a/examples/selective_build/test_selective_build.sh b/examples/selective_build/test_selective_build.sh index 5af3de5f3e5..36f6e3e59c1 100644 --- a/examples/selective_build/test_selective_build.sh +++ b/examples/selective_build/test_selective_build.sh @@ -66,7 +66,6 @@ test_buck2_select_ops_in_dict() { # select ops and their dtypes using the dictionary API. $BUCK run //examples/selective_build:selective_build_test \ --config=executorch.select_ops=dict \ - --config=executorch.dtype_selective_build_lib=//examples/selective_build:select_ops_in_dict_lib \ -- --model_path=./add_mul.pte echo "Removing add_mul.pte" diff --git a/shim/xplat/executorch/codegen/codegen.bzl b/shim/xplat/executorch/codegen/codegen.bzl index 4b69a2cf4a0..fd5ce84b789 100644 --- a/shim/xplat/executorch/codegen/codegen.bzl +++ b/shim/xplat/executorch/codegen/codegen.bzl @@ -42,6 +42,39 @@ CUSTOM_OPS_SCHEMA_REGISTRATION_SOURCES = [ "RegisterSchema.cpp", ] +ScalarType = enum( + "Byte", + "Char", + "Short", + "Int", + "Long", + "Half", + "Float", + "Double", + "ComplexHalf", + "ComplexFloat", + "ComplexDouble", + "Bool", + "QInt8", + "QUInt8", + "QInt32", + "BFloat16", + "QUInt4x2", + "QUInt2x4", + "Bits1x8", + "Bits2x4", + "Bits4x2", + "Bits8", + "Bits16", + "Float8_e5m2", + "Float8_e4m3fn", + "Float8_e5m2fnuz", + "Float8_e4m3fnuz", + "UInt16", + "UInt32", + "Uint64", +) + # Hide the dependency to caffe2 internally. def et_operator_library( name,