Generate ETKernelKeys from arg kernel specs Multiple ETKernelKeys are returned due to dtype permutations from utilizing type_alias_map (actualizing each potential type permutation as a KernelKey) Args: args: Mapping from argument name to kernel specs
(
args: Dict[str, Tuple[str, str]],
type_alias_map: Dict[str, List[str]], # TODO: Support unwrapped str val
dim_order_alias_map: Dict[str, List[int]],
)
| 60 | |
| 61 | @staticmethod |
| 62 | def gen_from_yaml( |
| 63 | args: Dict[str, Tuple[str, str]], |
| 64 | type_alias_map: Dict[str, List[str]], # TODO: Support unwrapped str val |
| 65 | dim_order_alias_map: Dict[str, List[int]], |
| 66 | ) -> List["ETKernelKey"]: |
| 67 | """Generate ETKernelKeys from arg kernel specs |
| 68 | Multiple ETKernelKeys are returned due to dtype permutations from utilizing |
| 69 | type_alias_map (actualizing each potential type permutation as a KernelKey) |
| 70 | |
| 71 | Args: |
| 72 | args: Mapping from argument name to kernel specs |
| 73 | Kernel specs are a tuple of (dtype, dim_order). |
| 74 | Currently tuple entries must be aliased via the alias map arguments |
| 75 | type_alias_map: Mapping from type alias to potential type enums |
| 76 | i.e { T0 : [Double, Int] } means T0 can be either Double or Int |
| 77 | Used for lookup by args |
| 78 | dim_order_alias_map: Mapping from alias to a list of dimension orders |
| 79 | Used for lookup by args |
| 80 | """ |
| 81 | # Cast to dim order to int |
| 82 | dim_order_alias_map = { |
| 83 | k: [int(alias) for alias in v] for k, v in dim_order_alias_map.items() |
| 84 | } |
| 85 | kernel_keys = [] |
| 86 | |
| 87 | # Get all used Dtype Alias |
| 88 | dtype_alias_used = set() |
| 89 | for type_alias, dim_order in args.values(): |
| 90 | # Enforce usage of alias initially |
| 91 | # TODO: Support inlined arguments |
| 92 | assert type_alias in type_alias_map, "Undefined type alias: " + str( |
| 93 | type_alias |
| 94 | ) |
| 95 | assert ( |
| 96 | dim_order in dim_order_alias_map |
| 97 | ), "Undefined dim_order alias: " + str(dim_order) |
| 98 | dtype_alias_used.add(type_alias) |
| 99 | |
| 100 | # Generate all permutations of dtype alias values |
| 101 | alias_dtypes = [ |
| 102 | [(alias, dtype) for dtype in type_alias_map[alias]] |
| 103 | for alias in dtype_alias_used |
| 104 | ] |
| 105 | alias_permutations = [ |
| 106 | dict(permutation) for permutation in list(itertools.product(*alias_dtypes)) |
| 107 | ] |
| 108 | |
| 109 | # Using each alias value permutation, generate kernel keys |
| 110 | op_arg_cache = {} |
| 111 | for permutation in alias_permutations: |
| 112 | arg_list = [] |
| 113 | for arg_name, arg_spec in args.items(): |
| 114 | dtype = permutation[arg_spec[0]] |
| 115 | dim_order = dim_order_alias_map[arg_spec[1]] # type: ignore[assignment] |
| 116 | if ( |
| 117 | cache_key := (arg_name, dtype, tuple(dim_order)) |
| 118 | ) not in op_arg_cache: |
| 119 | op_arg_cache[cache_key] = ETKernelKeyOpArgMeta(*cache_key) # type: ignore[arg-type] |