5 from collections
import OrderedDict
9 sys.path.append(path.dirname(path.abspath(__file__)))
14 import preprocess_declarations
15 import function_wrapper
17 from code_template
import CodeTemplate
27 parser = argparse.ArgumentParser(description=
'Generate ATen source files')
28 parser.add_argument(
'files', help=
'cwrap files', nargs=
'+')
33 help=
'path to source directory for ATen',
37 '--output-dependencies',
38 help=
'output a list of dependencies into the given file and exit')
40 '-d',
'--install_dir', help=
'output directory', default=
'ATen')
44 help=
'reinterpret CUDA as ROCm/HIP and adjust filepaths accordingly')
45 options = parser.parse_args()
46 gen_to_source = os.environ.get(
'GEN_TO_SOURCE')
48 core_install_dir = os.path.join(options.install_dir,
'core_tmp')
if options.install_dir
is not None else None 50 core_install_dir = os.path.join(options.source_path,
'core')
52 if options.install_dir
is not None and not os.path.exists(options.install_dir):
53 os.makedirs(options.install_dir)
54 if core_install_dir
is not None and not os.path.exists(core_install_dir):
55 os.makedirs(core_install_dir)
59 def __init__(self, install_dir=None):
60 self.
install_dir = install_dir
if install_dir
else options.install_dir
65 def will_write(self, filename):
66 filename =
'{}/{}'.format(self.
install_dir, filename)
68 raise Exception(
"'will_write' can only be called before " +
69 "the call to write_outputs, refactor so outputs are registered " +
70 "before running the generators")
71 self.filenames.add(filename)
73 def _write_if_changed(self, filename, contents):
75 with open(filename,
'r') as f: 76 old_contents = f.read() 79 if contents != old_contents:
80 with open(filename,
'w')
as f:
84 """Write a file containing the list of all outputs which are 85 generated by this script.""" 88 ''.join(name +
";" for name
in sorted(self.
filenames)))
91 def write(self, filename, s, env=None):
92 filename =
'{}/{}'.format(self.
install_dir, filename)
93 if isinstance(s, CodeTemplate):
94 assert env
is not None 95 env[
'generated_comment'] =
"@" +
"generated by aten/src/ATen/gen.py" 99 self.undeclared_files.append(filename)
101 self.filenames.remove(filename)
103 def check_all_files_written(self):
107 "in the list of outputs this script produces. " +
108 "use will_write to add them.")
110 raise Exception(
"Outputs declared with 'will_write' were " +
111 "never written: {}".format(self.
filenames))
114 TEMPLATE_PATH = options.source_path +
"/templates" 115 GENERATOR_DERIVED = CodeTemplate.from_file(
116 TEMPLATE_PATH +
"/GeneratorDerived.h")
117 TYPE_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeDerived.cpp")
118 SPARSE_TYPE_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/SparseTypeDerived.cpp")
119 TYPE_DERIVED_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeDerived.h")
120 TYPE_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/Type.h")
121 TYPE_EXTENDED_INTERFACE_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeExtendedInterface.h")
122 TYPE_DEFAULT_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeDefault.h")
123 TYPE_DEFAULT_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeDefault.cpp")
124 TYPE_EXTENSION_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeExtension.h")
125 TYPE_EXTENSION_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeExtension.cpp")
126 TYPE_EXTENSION_DERIVED_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeExtensionDerived.h")
127 TYPE_EXTENSION_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/TypeExtensionDerived.cpp")
129 LEGACY_TH_DISPATCHER_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/LegacyTHDispatcher.h")
130 LEGACY_TH_DISPATCHER_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/LegacyTHDispatcher.cpp")
131 LEGACY_TH_DISPATCHER_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/LegacyTHDispatcherDerived.cpp")
132 LEGACY_TH_DISPATCHER_DERIVED_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/LegacyTHDispatcherDerived.h")
134 REGISTER_CPU_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/RegisterCPU.h")
135 REGISTER_CPU_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/RegisterCPU.cpp")
137 REGISTER_CUDA_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/RegisterCUDA.h")
138 REGISTER_CUDA_CPP = CodeTemplate.from_file(TEMPLATE_PATH +
"/RegisterCUDA.cpp")
140 TENSOR_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/Tensor.h")
141 TENSOR_METHODS_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/TensorMethods.h")
143 FUNCTIONS_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/Functions.h")
144 LEGACY_TH_FUNCTIONS_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/LegacyTHFunctions.h")
146 NATIVE_FUNCTIONS_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/NativeFunctions.h")
148 EXTENSION_BACKEND_REGISTRATION_H = CodeTemplate.from_file(TEMPLATE_PATH +
"/ExtensionBackendRegistration.h")
151 context->registerType(Backend::${backend}, ScalarType::${scalar_type}, new ${type_name}()); 155 case Backend::${Backend}: 156 ${Type}Dispatch::register_function(schema, fn); 167 'th_generator':
'THGenerator * generator;',
173 'header':
'THC/THC.h' if not options.rocm
else 'THH/THH.h' 177 backends = [
'CPU',
'CUDA']
178 densities = [
'Dense',
'Sparse']
179 extension_backends = [
'MSNPU',
'XLA']
183 (
'Bool',
'uint8_t',
'BoolAccrealNotDefined',
'uint8_t',
False),
184 (
'Byte',
'uint8_t',
'Long',
'uint8_t',
False),
185 (
'Char',
'int8_t',
'Long',
'int8_t',
False),
186 (
'Double',
'double',
'Double',
'double',
True),
187 (
'Float',
'float',
'Double',
'float',
True),
188 (
'Int',
'int',
'Long',
'int32_t',
False),
189 (
'Long',
'int64_t',
'Long',
'int64_t',
False),
190 (
'Short',
'int16_t',
'Long',
'int16_t',
False),
191 (
'Half',
'Half',
'Double',
'at::Half',
True),
196 'cpu_type_registrations': [],
197 'cpu_type_headers': [],
198 'cuda_type_registrations': [],
199 'cuda_type_headers': [],
200 'pure_virtual_type_method_declarations': [],
201 'pure_virtual_extended_type_method_declarations': [],
202 'type_method_declarations': [],
203 'type_method_definitions': [],
204 'tensor_method_declarations': [],
205 'tensor_method_definitions': [],
206 'function_declarations': [],
207 'function_definitions': [],
209 'native_function_declarations': [],
210 'extension_backend_headers': [],
211 'extension_backend_register_switches': [],
215 def dict_representer(dumper, data):
216 return dumper.represent_dict(data.items())
219 def postprocess_output_declarations(output_declarations):
221 for decl
in output_declarations:
222 has_named_ret =
False 223 for n, ret
in enumerate(decl.returns):
224 if 'name' not in ret:
225 assert not has_named_ret
228 elif len(decl.returns) == 1:
231 ret[
'name'] =
'out' + str(n)
235 def remove_key_if_none(dictionary, key):
236 if key
in dictionary.keys()
and dictionary[key]
is None:
240 return [remove_key_if_none(decl._asdict(),
'buffers')
241 for decl
in output_declarations]
244 def format_yaml(data):
245 if options.output_dependencies:
248 noalias_dumper = yaml.dumper.SafeDumper
249 noalias_dumper.ignore_aliases =
lambda self, data:
True 251 noalias_dumper.add_representer(OrderedDict, dict_representer)
255 return yaml.dump(data, default_flow_style=
False, Dumper=noalias_dumper, width=float(
'Inf'))
258 def generate_storage_type_and_tensor(backend, density, scalar_type, declarations):
259 scalar_name, c_type, accreal, th_scalar_type, is_floating_type = scalar_type
261 density_tag =
'Sparse' if density ==
'Sparse' else '' 262 env[
'Density'] = density
263 env[
'ScalarName'] = scalar_name
264 env[
'ScalarType'] = c_type
265 env[
'THScalarType'] = th_scalar_type
266 env[
'AccScalarName'] = accreal
267 env[
'isFloatingType'] = is_floating_type
268 env[
'isIntegralType'] =
not is_floating_type
269 env[
'Type'] =
"{}{}{}Type".format(density_tag, backend, scalar_name)
270 env[
'DenseTensor'] =
"{}{}Tensor".format(backend, scalar_name)
271 env[
'Backend'] = density_tag + backend
272 env[
'DenseBackend'] = backend
273 env[
'storage_tensor_headers'] = []
274 if density !=
'Sparse':
275 env[
'storage_tensor_headers'] = [
'#include <c10/core/TensorImpl.h>']
278 tag = density_tag + backend + scalar_name
279 env[
'TypeID'] =
'TypeID::' + tag
280 top_env[
'type_ids'].append(tag +
',')
282 if backend ==
'CUDA':
283 env[
'extra_cuda_headers'] = []
284 env[
'extra_cuda_headers'].append(
'#include <ATen/DeviceGuard.h>')
286 env[
'th_headers'] = [
287 '#include <THH/THH.h>',
288 '#include <THH/THHTensor.hpp>',
289 '#include <THHUNN/THHUNN.h>',
291 '#undef THCIndexTensor_',
293 env[
'extra_cuda_headers'].append(
'#include <ATen/hip/ATenHIPGeneral.h>')
294 env[
'extra_cuda_headers'].append(
'#include <ATen/hip/HIPDevice.h>')
295 env[
'extra_cuda_headers'].append(
'#include <ATen/hip/HIPTypeDefault.h>')
297 env[
'th_headers'] = [
298 '#include <THC/THC.h>',
299 '#include <THC/THCTensor.hpp>',
300 '#include <THCUNN/THCUNN.h>',
302 '#undef THCIndexTensor_',
304 env[
'extra_cuda_headers'].append(
'#include <ATen/cuda/ATenCUDAGeneral.h>')
305 env[
'extra_cuda_headers'].append(
'#include <ATen/cuda/CUDADevice.h>')
306 env[
'extra_cuda_headers'].append(
'#include <ATen/cuda/CUDATypeDefault.h>')
307 sname =
'' if scalar_name ==
"Float" else scalar_name
308 env[
'THType'] =
'Cuda{}'.format(sname)
309 env[
'THStorage'] =
'THCuda{}Storage'.format(sname)
310 env[
'THTensor'] =
'THCuda{}Tensor'.format(sname)
311 env[
'THIndexTensor'] =
'THCudaLongTensor' 312 env[
'state'] = [
'globalContext().getTHCState()']
313 env[
'isCUDA'] =
'true' 314 env[
'storage_device'] =
'return storage->device;' 315 env[
'Generator'] =
'CUDAGenerator' 317 env[
'th_headers'] = [
318 '#include <TH/TH.h>',
319 '#include <TH/THTensor.hpp>',
320 '#include <THNN/THNN.h>',
323 env[
'extra_cuda_headers'] = []
324 env[
'THType'] = scalar_name
325 env[
'THStorage'] =
"TH{}Storage".format(scalar_name)
326 env[
'THTensor'] =
'TH{}Tensor'.format(scalar_name)
327 env[
'THIndexTensor'] =
'THLongTensor' 329 env[
'isCUDA'] =
'false' 330 env[
'storage_device'] =
'throw std::runtime_error("CPU storage has no device");' 331 env[
'Generator'] =
'CPUGenerator' 332 env[
'AS_REAL'] = env[
'ScalarType']
333 if scalar_name ==
"Half":
334 env[
'SparseTensor'] =
'Tensor' 335 if backend ==
"CUDA":
336 env[
'AS_REAL'] =
'convert<at::Half,double>' 338 declarations, definitions = function_wrapper.create_derived(
340 env[
'type_derived_method_declarations'] = declarations
341 env[
'type_derived_method_definitions'] = definitions
344 if env[
'DenseBackend'] ==
'CUDA':
345 fm = cuda_file_manager
347 if density !=
'Sparse':
348 fm.write(env[
'Type'] +
".cpp", TYPE_DERIVED_CPP, env)
350 fm.write(env[
'Type'] +
".cpp", SPARSE_TYPE_DERIVED_CPP, env)
351 fm.write(env[
'Type'] +
".h", TYPE_DERIVED_H, env)
353 type_register = TYPE_REGISTER.substitute(backend=env[
'Backend'], scalar_type=scalar_name, type_name=env[
'Type'])
354 if env[
'DenseBackend'] ==
'CPU':
355 top_env[
'cpu_type_registrations'].append(type_register)
356 top_env[
'cpu_type_headers'].append(
357 '#include "ATen/{}.h"'.format(env[
'Type']))
359 assert env[
'DenseBackend'] ==
'CUDA' 360 top_env[
'cuda_type_registrations'].append(type_register)
361 top_env[
'cuda_type_headers'].append(
362 '#include "ATen/{}.h"'.format(env[
'Type']))
365 def generate_type_extension_backend(backend, declarations):
367 env[
'Type'] =
"{}Type".format(backend)
368 env[
'Backend'] = backend
369 env[
'DeviceType'] = backend
371 declarations, definitions = function_wrapper.create_extension_backend(
373 env[
'type_method_declarations'] = declarations
374 env[
'type_method_definitions'] = definitions
376 file_manager.write(env[
'Type'] +
".cpp", TYPE_EXTENSION_CPP, env)
377 file_manager.write(env[
'Type'] +
".h", TYPE_EXTENSION_H, env)
379 extension_backend_register_switch = EXTENSION_BACKEND_REGISTER_SWITCH.substitute(env)
380 top_env[
'extension_backend_register_switches'].append(extension_backend_register_switch)
381 top_env[
'extension_backend_headers'].append(
382 '#include <ATen/{}.h>'.format(env[
'Type']))
385 def generate_type_extension_backend_derived_types(backend):
387 env[
'Backend'] = backend
388 for scalar_name, c_type, _, _, _
in scalar_types:
389 env[
'Type'] =
"{}{}Type".format(backend, scalar_name)
390 env[
'ScalarName'] = scalar_name
391 env[
'ScalarType'] = c_type
392 env[
'TypeID'] =
'TypeID::' + backend + scalar_name
393 top_env[
'type_ids'].append(backend + scalar_name +
',')
395 type_register = TYPE_REGISTER.substitute(backend=env[
'Backend'], scalar_type=scalar_name, type_name=env[
'Type'])
396 top_env[
'cpu_type_registrations'].append(type_register)
397 file_manager.write(env[
'Type'] +
".cpp", TYPE_EXTENSION_DERIVED_CPP, env)
398 file_manager.write(env[
'Type'] +
".h", TYPE_EXTENSION_DERIVED_H, env)
400 top_env[
'cpu_type_headers'].append(
'#include "ATen/{}.h"'.format(env[
'Type']))
403 def generate_legacy_th_dispatcher(backend, density, scalar_type, declarations):
404 assert density !=
'Sparse' 405 scalar_name, c_type, accreal, th_scalar_type, is_floating_type = scalar_type
407 env[
'Backend'] = backend
408 env[
'Dispatcher'] =
"LegacyTH{}{}Dispatcher".format(backend, scalar_name)
411 if backend ==
'CUDA':
412 fm = cuda_file_manager
414 fm.write(env[
'Dispatcher'] +
".cpp", LEGACY_TH_DISPATCHER_DERIVED_CPP, env)
415 fm.write(env[
'Dispatcher'] +
".h", LEGACY_TH_DISPATCHER_DERIVED_H, env)
421 for backend
in backends:
422 for density
in densities:
423 for scalar_type
in scalar_types:
424 if density ==
'Sparse' and scalar_type[0] ==
'Half':
427 yield (backend, density, scalar_type)
434 def declare_outputs():
435 core_files = [
'Type.h',
'Tensor.h',
'TensorMethods.h']
437 core_file_manager.will_write(f)
438 files = [
'Declarations.yaml',
'TypeExtendedInterface.h',
'TypeDefault.cpp',
'TypeDefault.h',
439 'LegacyTHDispatcher.h',
'LegacyTHDispatcher.cpp',
'LegacyTHFunctions.h',
440 'Functions.h',
'NativeFunctions.h',
'RegisterCPU.cpp',
'RegisterCPU.h',
'ExtensionBackendRegistration.h']
442 file_manager.will_write(f)
443 cuda_files = [
'RegisterCUDA.cpp',
'RegisterCUDA.h']
445 cuda_file_manager.will_write(f)
446 for fname
in sorted(generators.keys()):
448 if generators[fname][
'name'] ==
'CUDA':
449 fm = cuda_file_manager
451 for backend, density, scalar_type
in iterate_types():
452 scalar_name = scalar_type[0]
453 full_backend =
"Sparse" + backend
if density ==
"Sparse" else backend
455 if backend ==
'CUDA':
456 fm = cuda_file_manager
457 for kind
in [
"Type"]:
458 if kind !=
'Type' and density ==
"Sparse":
461 fm.will_write(
"{}{}{}.h".format(full_backend, scalar_name, kind))
462 fm.will_write(
"{}{}{}.cpp".format(full_backend, scalar_name, kind))
464 if density !=
'Sparse':
465 fm.will_write(
"{}{}{}{}.h".format(
'LegacyTH', full_backend, scalar_name,
'Dispatcher'))
466 fm.will_write(
"{}{}{}{}.cpp".format(
'LegacyTH', full_backend, scalar_name,
'Dispatcher'))
467 for backend
in extension_backends:
468 file_manager.will_write(
"{}Type.h".format(backend))
469 file_manager.will_write(
"{}Type.cpp".format(backend))
470 for scalar_type
in scalar_types:
471 scalar_name = scalar_type[0]
472 file_manager.will_write(
"{}{}Type.h".format(backend, scalar_name))
473 file_manager.will_write(
"{}{}Type.cpp".format(backend, scalar_name))
476 def filter_by_extension(files, *extensions):
479 for extension
in extensions:
480 if file.endswith(extension):
481 filtered_files.append(file)
482 return filtered_files
487 def cmpfiles_with_eol_normalization(a, b, names):
488 results = ([], [], [])
491 with open(os.path.join(a, x))
as f:
492 ax = f.read().replace(
'\r\n',
'\n').replace(
'\r',
'\n')
493 with open(os.path.join(b, x))
as f:
494 bx = f.read().replace(
'\r\n',
'\n').replace(
'\r',
'\n')
504 def generate_outputs():
505 cwrap_files = filter_by_extension(options.files,
'.cwrap')
506 nn_files = filter_by_extension(options.files,
'nn.yaml',
'.h')
507 native_files = filter_by_extension(options.files,
'native_functions.yaml')
510 for file
in cwrap_files
516 for fname, env
in generators.items():
518 if env[
'name'] ==
'CUDA':
519 fm = cuda_file_manager
520 fm.write(fname, GENERATOR_DERIVED, env)
525 output_declarations = function_wrapper.create_generic(top_env, declarations)
526 output_declarations = postprocess_output_declarations(output_declarations)
527 file_manager.write(
"Declarations.yaml", format_yaml(output_declarations))
529 for backend, density, scalar_type
in iterate_types():
530 generate_storage_type_and_tensor(backend, density, scalar_type, declarations)
531 for backend
in extension_backends:
532 generate_type_extension_backend(backend, declarations)
533 generate_type_extension_backend_derived_types(backend)
535 for backend, density, scalar_type
in iterate_types():
536 if density !=
'Sparse':
537 generate_legacy_th_dispatcher(backend, density, scalar_type, [])
541 'Tensor.h': TENSOR_H,
542 'TensorMethods.h': TENSOR_METHODS_H
545 for core_file, core_template_file
in core_files.items():
546 core_file_manager.write(core_file, core_template_file, top_env)
548 file_manager.write(
'TypeExtendedInterface.h', TYPE_EXTENDED_INTERFACE_H, top_env)
549 file_manager.write(
'TypeDefault.h', TYPE_DEFAULT_H, top_env)
550 file_manager.write(
'TypeDefault.cpp', TYPE_DEFAULT_CPP, top_env)
552 file_manager.write(
'LegacyTHDispatcher.h', LEGACY_TH_DISPATCHER_H, top_env)
553 file_manager.write(
'LegacyTHDispatcher.cpp', LEGACY_TH_DISPATCHER_CPP, top_env)
555 file_manager.write(
'RegisterCPU.h', REGISTER_CPU_H, top_env)
556 file_manager.write(
'RegisterCPU.cpp', REGISTER_CPU_CPP, top_env)
558 cuda_file_manager.write(
'RegisterCUDA.h', REGISTER_CUDA_H, top_env)
559 cuda_file_manager.write(
'RegisterCUDA.cpp', REGISTER_CUDA_CPP, top_env)
561 file_manager.write(
'Functions.h', FUNCTIONS_H, top_env)
562 file_manager.write(
'LegacyTHFunctions.h', LEGACY_TH_FUNCTIONS_H, top_env)
564 file_manager.write(
'NativeFunctions.h', NATIVE_FUNCTIONS_H, top_env)
566 file_manager.write(
'ExtensionBackendRegistration.h', EXTENSION_BACKEND_REGISTRATION_H, top_env)
568 file_manager.check_all_files_written()
569 cuda_file_manager.check_all_files_written()
572 core_source_path = os.path.join(options.source_path,
'core')
573 match, mismatch, errors = cmpfiles_with_eol_normalization(core_install_dir, core_source_path, core_files.keys())
575 raise RuntimeError(
"Error while trying to compare source and generated files for {}. " 576 "Source directory: {}. Generated directory: {}." 577 .format(errors, core_source_path, core_install_dir))
579 file_component =
'{}'.format(
','.join(mismatch))
580 if len(mismatch) > 1:
581 file_component =
'{' + file_component +
'}' 582 update_cmd =
"cp {}/{} {}".format(core_install_dir, file_component, core_source_path)
583 raise RuntimeError(
"Source files: {} did not match generated files. To update the source files, " 584 "set environment variable GEN_TO_SOURCE or run \"{}\"".format(mismatch, update_cmd))
587 if options.output_dependencies
is not None:
588 file_manager.write_outputs(options.output_dependencies)
589 core_file_manager.write_outputs(options.output_dependencies +
"-core")
590 cuda_file_manager.write_outputs(options.output_dependencies +
"-cuda")
def _write_if_changed(self, filename, contents)
def write_outputs(self, filename)