| 50 | |
| 51 | @dataclass(frozen=True) |
| 52 | class ETKernelKey: |
| 53 | # Field undefined is default = True |
| 54 | arg_meta: Tuple[ETKernelKeyOpArgMeta, ...] = () |
| 55 | |
| 56 | # Indicator for this kernel being used as a catch all |
| 57 | default: bool = False |
| 58 | |
| 59 | version: int = KERNEL_KEY_VERSION |
| 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 |
no outgoing calls
searching dependent graphs…