(self)
| 155 | self._write_def("MGB_OPR_REGISTRY_CALLER_SPECIALIZE", defs) |
| 156 | |
| 157 | def _write_elemwise_modes(self): |
| 158 | with tempfile.NamedTemporaryFile() as ftmp: |
| 159 | fpath = os.path.realpath(ftmp.name) |
| 160 | subprocess.check_call( |
| 161 | [ |
| 162 | "./dnn/scripts/gen_param_defs.py", |
| 163 | "--write-enum-items", |
| 164 | "Elemwise:Mode", |
| 165 | "./dnn/scripts/opr_param_defs.py", |
| 166 | fpath, |
| 167 | ], |
| 168 | cwd=self.get_megengine_root(), |
| 169 | ) |
| 170 | |
| 171 | with open(fpath) as fin: |
| 172 | mode_list = [i.strip() for i in fin] |
| 173 | |
| 174 | all_elemwise_modes = set() |
| 175 | for i in mode_list: |
| 176 | i_type = i.replace(" ", "").replace("=", " ").split()[0] |
| 177 | i_id = i.replace(" ", "").replace("=", " ").split()[1] |
| 178 | all_elemwise_modes.add(i_id) |
| 179 | |
| 180 | if i_id in self._elemwise_modes: |
| 181 | content = "_cb({})".format(i_type) |
| 182 | else: |
| 183 | content = "" |
| 184 | self._write_def( |
| 185 | "_MEGDNN_ELEMWISE_MODE_ENABLE_IMPL_{}(_cb)".format(i_type), content, |
| 186 | ) |
| 187 | |
| 188 | # write end of elemwise macro |
| 189 | self._write_def( |
| 190 | "MEGDNN_ELEMWISE_MODE_ENABLE(_mode, _cb)", |
| 191 | "_MEGDNN_ELEMWISE_MODE_ENABLE_IMPL_##_mode(_cb)", |
| 192 | ) |
| 193 | # finally check all self._elemwise_modes is in all_elemwise_modes |
| 194 | for i in self._elemwise_modes: |
| 195 | assert ( |
| 196 | i in all_elemwise_modes |
| 197 | ), "code issue happened, can not find elemwise mode: {} in {}".format( |
| 198 | i, all_elemwise_modes |
| 199 | ) |
| 200 | |
| 201 | def _write_dtype(self): |
| 202 | if "Float16" not in self._dtypes: |
no test coverage detected