def generate_storage_type_and_tensor(backend, density, scalar_type,
declarations):
scalar_name, c_type, accreal, th_scalar_type = scalar_type
env = {}
density_tag = 'Sparse' if density == 'Sparse' else ''
th_density_tag = 'S' if density == 'Sparse' else ''
env['Density'] = density
env['ScalarName'] = scalar_name
env['ScalarType'] = c_type
env['THScalarType'] = th_scalar_type
env['AccScalarName'] = accreal
env['Storage'] = "{}{}Storage".format(backend, scalar_name)
env['Type'] = "{}{}{}Type".format(density_tag, backend, scalar_name)
env['Tensor'] = "{}{}{}Tensor".format(density_tag, backend, scalar_name)
env['SparseTensor'] = "Sparse{}{}Tensor".format(backend, scalar_name)
env['Backend'] = density_tag + backend
# used for generating switch logic for external functions
tag = density_tag + backend + scalar_name
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
if backend == 'CUDA':
env['th_headers'] = [
'#include <THC/THC.h>', '#include <THCUNN/THCUNN.h>',
'#undef THNN_', '#undef THCIndexTensor_'
]
# if density == 'Sparse':
env['th_headers'] += [
'#include <THCS/THCS.h>', '#undef THCIndexTensor_'
]
sname = '' if scalar_name == "Float" else scalar_name
env['THType'] = 'Cuda{}'.format(sname)
env['THStorage'] = 'THCuda{}Storage'.format(sname)
if density == 'Dense':
env['THTensor'] = 'THCuda{}Tensor'.format(sname)
else:
env['THTensor'] = 'THCS{}Tensor'.format(scalar_name)
env['THIndexTensor'] = 'THCudaLongTensor'
env['state'] = ['context->thc_state']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
env['Generator'] = 'CUDAGenerator'
else:
env['th_headers'] = [
'#include <TH/TH.h>', '#include <THNN/THNN.h>', '#undef THNN_'
]
# if density == 'Sparse':
env['th_headers'].append('#include <THS/THS.h>')
env['THType'] = scalar_name
env['THStorage'] = "TH{}Storage".format(scalar_name)
env['THTensor'] = 'TH{}{}Tensor'.format(th_density_tag, scalar_name)
env['THIndexTensor'] = 'THLongTensor'
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['Generator'] = 'CPUGenerator'
env['AS_REAL'] = env['ScalarType']
if scalar_name == "Half":
env['SparseTensor'] = 'Tensor'
if backend == "CUDA":
env['to_th_type'] = 'HalfFix<__half,Half>'
env['to_at_type'] = 'HalfFix<Half,__half>'
env['AS_REAL'] = 'convert<half,double>'
env['THScalarType'] = 'half'
else:
env['to_th_type'] = 'HalfFix<THHalf,Half>'
env['to_at_type'] = 'HalfFix<Half,THHalf>'
elif scalar_name == 'Long':
env['to_th_type'] = 'long'
env['to_at_type'] = 'int64_t'
else:
env['to_th_type'] = ''
env['to_at_type'] = ''
declarations, definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
if density != 'Sparse':
# there are no special storage types for Sparse, they are composed
# of Dense tensors
write(env['Storage'] + ".cpp", STORAGE_DERIVED_CPP.substitute(env))
write(env['Storage'] + ".h", STORAGE_DERIVED_H.substitute(env))
env['TensorDenseOrSparse'] = TENSOR_DENSE_CPP.substitute(env)
env['THTensor_nDimension'] = 'tensor->nDimension'
else:
env['TensorDenseOrSparse'] = TENSOR_SPARSE_CPP.substitute(env)
env['THTensor_nDimension'] = 'tensor->nDimensionI + tensor->nDimensionV'
write(env['Type'] + ".cpp", TYPE_DERIVED_CPP.substitute(env))
write(env['Type'] + ".h", TYPE_DERIVED_H.substitute(env))
write(env['Tensor'] + ".cpp", TENSOR_DERIVED_CPP.substitute(env))
write(env['Tensor'] + ".h", TENSOR_DERIVED_H.substitute(env))
type_register = (
('context->type_registry[static_cast<int>(Backend::{})]' +
'[static_cast<int>(ScalarType::{})].reset(new {}(context));').format(
env['Backend'], scalar_name, env['Type']))
top_env['type_registrations'].append(type_register)
top_env['type_headers'].append('#include "ATen/{}.h"'.format(env['Type']))
return env
def generate_storage_type_and_tensor(backend, density, scalar_type, declarations):
scalar_name, c_type, accreal, th_scalar_type, is_floating_type = scalar_type
env = {}
density_tag = 'Sparse' if density == 'Sparse' else ''
th_density_tag = 'S' if density == 'Sparse' else ''
env['Density'] = density
env['ScalarName'] = scalar_name
env['ScalarType'] = c_type
env['THScalarType'] = th_scalar_type
env['AccScalarName'] = accreal
env['isFloatingType'] = is_floating_type
env['isIntegralType'] = not is_floating_type
env['Storage'] = "{}{}Storage".format(backend, scalar_name)
env['Type'] = "{}{}{}Type".format(density_tag, backend, scalar_name)
env['Tensor'] = "{}{}{}Tensor".format(density_tag, backend, scalar_name)
env['DenseTensor'] = "{}{}Tensor".format(backend, scalar_name)
env['SparseTensor'] = "Sparse{}{}Tensor".format(backend, scalar_name)
env['Backend'] = density_tag + backend
env['DenseBackend'] = backend
# used for generating switch logic for external functions
tag = density_tag + backend + scalar_name
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
if backend == 'CUDA':
env['th_headers'] = [
'#include <THC/THC.h>',
'#include <THC/THCTensor.hpp>',
'#include <THCUNN/THCUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
'#include <THCS/THCS.h>',
'#include <THCS/THCSTensor.hpp>',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'] = ['#include <ATen/cuda/CUDAHalf.cuh>']
sname = '' if scalar_name == "Float" else scalar_name
env['THType'] = 'Cuda{}'.format(sname)
env['THStorage'] = 'THCuda{}Storage'.format(sname)
if density == 'Dense':
env['THTensor'] = 'THCuda{}Tensor'.format(sname)
else:
env['THTensor'] = 'THCS{}Tensor'.format(scalar_name)
env['THIndexTensor'] = 'THCudaLongTensor'
env['state'] = ['context->thc_state']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
env['Generator'] = 'CUDAGenerator'
else:
env['th_headers'] = [
'#include <TH/TH.h>',
'#include <TH/THTensor.hpp>',
'#include <THNN/THNN.h>',
'#undef THNN_',
'#include <THS/THS.h>',
'#include <THS/THSTensor.hpp>',
]
env['extra_cuda_headers'] = []
env['THType'] = scalar_name
env['THStorage'] = "TH{}Storage".format(scalar_name)
env['THTensor'] = 'TH{}{}Tensor'.format(th_density_tag, scalar_name)
env['THIndexTensor'] = 'THLongTensor'
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['Generator'] = 'CPUGenerator'
env['AS_REAL'] = env['ScalarType']
if scalar_name == "Half":
env['SparseTensor'] = 'Tensor'
if backend == "CUDA":
env['to_th_type'] = 'HalfFix<__half,Half>'
env['to_at_type'] = 'HalfFix<Half,__half>'
env['AS_REAL'] = 'convert<half,double>'
env['THScalarType'] = 'half'
else:
env['to_th_type'] = 'HalfFix<THHalf,Half>'
env['to_at_type'] = 'HalfFix<Half,THHalf>'
elif scalar_name == 'Long':
env['to_th_type'] = 'long'
env['to_at_type'] = 'int64_t'
else:
env['to_th_type'] = ''
env['to_at_type'] = ''
declarations, definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
if density != 'Sparse':
# there are no special storage types for Sparse, they are composed
# of Dense tensors
file_manager.write(env['Storage'] + ".cpp", STORAGE_DERIVED_CPP.substitute(env))
file_manager.write(env['Storage'] + ".h", STORAGE_DERIVED_H.substitute(env))
env['TensorDenseOrSparse'] = TENSOR_DENSE_CPP.substitute(env)
env['THTensor_nDimension'] = 'tensor->nDimension'
else:
env['TensorDenseOrSparse'] = TENSOR_SPARSE_CPP.substitute(env)
env['THTensor_nDimension'] = 'tensor->nDimensionI + tensor->nDimensionV'
file_manager.write(env['Type'] + ".cpp", TYPE_DERIVED_CPP.substitute(env))
file_manager.write(env['Type'] + ".h", TYPE_DERIVED_H.substitute(env))
file_manager.write(env['Tensor'] + ".cpp", TENSOR_DERIVED_CPP.substitute(env))
file_manager.write(env['Tensor'] + ".h", TENSOR_DERIVED_H.substitute(env))
type_register = (('context->type_registry[static_cast<int>(Backend::{})]' +
'[static_cast<int>(ScalarType::{})].reset(new {}(context));')
.format(env['Backend'], scalar_name, env['Type']))
top_env['type_registrations'].append(type_register)
top_env['type_headers'].append(
'#include "ATen/{}.h"'.format(env['Type']))
return env
def generate_storage_type_and_tensor(backend, density, scalar_type, declarations):
scalar_name, c_type, accreal, is_floating_type = scalar_type
env = {}
density_tag = density if density != 'Dense' else ''
env['Density'] = density
env['ScalarName'] = scalar_name
env['ScalarType'] = c_type
env['AccScalarName'] = accreal
env['isFloatingType'] = is_floating_type
env['isIntegralType'] = not is_floating_type
env['Type'] = "{}{}{}Type".format(density_tag, backend, scalar_name)
env['DenseTensor'] = "{}{}Tensor".format(backend, scalar_name)
env['Backend'] = density_tag + backend
env['DenseBackend'] = backend
env['storage_tensor_headers'] = []
if density != 'Sparse':
env['storage_tensor_headers'] = ['#include <c10/core/TensorImpl.h>']
# used for generating switch logic for external functions
tag = density_tag + backend + scalar_name
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
if backend == 'CUDA':
env['extra_cuda_headers'] = []
env['extra_cuda_headers'].append('#include <ATen/DeviceGuard.h>')
if options.rocm:
env['th_headers'] = [
'#include <THH/THH.h>',
'#include <THH/THHTensor.hpp>',
'#include <THHUNN/THHUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'].append('#include <ATen/hip/ATenHIPGeneral.h>')
env['extra_cuda_headers'].append('#include <ATen/hip/HIPDevice.h>')
env['extra_cuda_headers'].append('#include <ATen/hip/HIPTypeDefault.h>')
else:
env['th_headers'] = [
'#include <THC/THC.h>',
'#include <THC/THCTensor.hpp>',
'#include <THCUNN/THCUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'].append('#include <ATen/cuda/ATenCUDAGeneral.h>')
env['extra_cuda_headers'].append('#include <ATen/cuda/CUDADevice.h>')
env['extra_cuda_headers'].append('#include <ATen/cuda/CUDATypeDefault.h>')
sname = '' if scalar_name == "Float" else scalar_name
env['THType'] = 'Cuda{}'.format(sname)
env['THStorage'] = 'THCuda{}Storage'.format(sname)
env['THTensor'] = 'THCuda{}Tensor'.format(sname)
env['THIndexTensor'] = 'THCudaLongTensor'
env['state'] = ['globalContext().getTHCState()']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
env['Generator'] = 'CUDAGenerator'
else:
env['th_headers'] = [
'#include <TH/TH.h>',
'#include <TH/THTensor.hpp>',
'#include <THNN/THNN.h>',
'#undef THNN_',
]
env['extra_cuda_headers'] = []
env['THType'] = scalar_name
env['THStorage'] = "TH{}Storage".format(scalar_name)
env['THTensor'] = 'TH{}Tensor'.format(scalar_name)
env['THIndexTensor'] = 'THLongTensor'
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['Generator'] = 'CPUGenerator'
env['AS_REAL'] = env['ScalarType']
if scalar_name == "Half":
env['SparseTensor'] = 'Tensor'
if backend == "CUDA":
env['AS_REAL'] = 'convert<at::Half,double>'
declarations, definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
fm = file_manager
if env['DenseBackend'] == 'CUDA':
fm = cuda_file_manager
if density != 'Sparse':
fm.write(env['Type'] + ".cpp", TYPE_DERIVED_CPP, env)
else:
fm.write(env['Type'] + ".cpp", SPARSE_TYPE_DERIVED_CPP, env)
fm.write(env['Type'] + ".h", TYPE_DERIVED_H, env)
type_register = TYPE_REGISTER.substitute(backend=env['Backend'], scalar_type=scalar_name, type_name=env['Type'])
if env['DenseBackend'] == 'CPU':
top_env['cpu_type_registrations'].append(type_register)
top_env['cpu_type_headers'].append(
'#include "ATen/{}.h"'.format(env['Type']))
else:
assert env['DenseBackend'] == 'CUDA'
top_env['cuda_type_registrations'].append(type_register)
top_env['cuda_type_headers'].append(
'#include "ATen/{}.h"'.format(env['Type']))
def generate_storage_type_and_tensor(backend, density, declarations):
env = {}
density_tag = density if density != 'Dense' else ''
env['Density'] = density
env['Type'] = "{}{}Type".format(density_tag, backend)
env['DeviceType'] = backend_to_devicetype(backend)
env['Backend'] = density_tag + backend
env['storage_tensor_headers'] = []
if density != 'Sparse':
env['storage_tensor_headers'] = ['#include <c10/core/TensorImpl.h>']
# used for generating switch logic for external functions
tag = density_tag + backend
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
env['legacy_th_headers'] = []
if backend == 'CUDA':
env['extra_cuda_headers'] = []
env['extra_cuda_headers'].append('#include <ATen/DeviceGuard.h>')
if options.rocm:
env['th_headers'] = [
'#include <THH/THH.h>',
'#include <THH/THHTensor.hpp>',
'#include <THHUNN/THHUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'].append(
'#include <ATen/hip/ATenHIPGeneral.h>')
env['extra_cuda_headers'].append('#include <ATen/hip/HIPDevice.h>')
env['extra_cuda_headers'].append(
'#include <ATen/hip/HIPContext.h>')
else:
env['th_headers'] = [
'#include <THC/THC.h>',
'#include <THC/THCTensor.hpp>',
'#include <THCUNN/THCUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'].append(
'#include <ATen/cuda/ATenCUDAGeneral.h>')
env['extra_cuda_headers'].append(
'#include <ATen/cuda/CUDADevice.h>')
env['extra_cuda_headers'].append(
'#include <ATen/cuda/CUDAContext.h>')
env['state'] = ['globalContext().getTHCState()']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
env['Generator'] = 'CUDAGenerator'
env['allocator'] = 'at::cuda::getCUDADeviceAllocator()'
else:
env['th_headers'] = [
'#include <TH/TH.h>',
'#include <TH/THTensor.hpp>',
'#include <THNN/THNN.h>',
'#undef THNN_',
]
env['extra_cuda_headers'] = []
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['Generator'] = 'CPUGenerator'
env['allocator'] = 'getCPUAllocator()'
declarations, definitions, registrations, th_declarations, th_definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
env['function_registrations'] = registrations
env['legacy_th_declarations'] = th_declarations
env['legacy_th_definitions'] = th_definitions
fm = file_manager
if env['DeviceType'] == 'CUDA':
fm = cuda_file_manager
if env['Backend'] == 'CPU' or env['Backend'] == 'CUDA':
env['namespace'] = env['Backend'].lower()
env['legacy_th_headers'].append('#include <ATen/LegacyTHFunctions' +
env['Backend'] + ".h>")
fm.write('LegacyTHFunctions' + env['Backend'] + ".h",
LEGACY_TH_FUNCTIONS_H, env)
fm.write('LegacyTHFunctions' + env['Backend'] + ".cpp",
LEGACY_TH_FUNCTIONS_CPP, env)
if density != 'Sparse':
fm.write(env['Type'] + ".cpp", TYPE_DERIVED_CPP, env)
else:
fm.write(env['Type'] + ".cpp", SPARSE_TYPE_DERIVED_CPP, env)
fm.write(env['Type'] + ".h", TYPE_DERIVED_H, env)
if env['DeviceType'] == 'CPU':
top_env['cpu_type_headers'].append('#include "ATen/{}.h"'.format(
env['Type']))
else:
assert env['DeviceType'] == 'CUDA'
top_env['cuda_type_headers'].append('#include "ATen/{}.h"'.format(
env['Type']))
def generate_storage_type_and_tensor(backend, density, scalar_type,
declarations):
scalar_name, c_type, accreal, th_scalar_type, is_floating_type = scalar_type
env = {}
density_tag = 'Sparse' if density == 'Sparse' else ''
env['Density'] = density
env['ScalarName'] = scalar_name
env['ScalarType'] = c_type
env['THScalarType'] = th_scalar_type
env['AccScalarName'] = accreal
env['isFloatingType'] = is_floating_type
env['isIntegralType'] = not is_floating_type
if density == 'Dense':
env['Storage'] = "{}{}Storage".format(backend, scalar_name)
env['Tensor'] = "{}{}{}Tensor".format(density_tag, backend,
scalar_name)
env['Type'] = "{}{}{}Type".format(density_tag, backend, scalar_name)
env['DenseTensor'] = "{}{}Tensor".format(backend, scalar_name)
env['Backend'] = density_tag + backend
env['DenseBackend'] = backend
env['storage_tensor_headers'] = []
if density != 'Sparse':
env['storage_tensor_headers'] = [
'#include "ATen/{}.h"'.format(env['Storage']),
'#include "ATen/{}.h"'.format(env['Tensor']),
'#include "ATen/{}ByteTensor.h"'.format(env['Backend']),
'#include "ATen/{}IntTensor.h"'.format(env['Backend']),
'#include "ATen/{}LongTensor.h"'.format(env['Backend']),
]
# used for generating switch logic for external functions
tag = density_tag + backend + scalar_name
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
if backend == 'CUDA':
env['th_headers'] = [
'#include <THC/THC.h>',
'#include <THC/THCTensor.hpp>',
'#include <THCUNN/THCUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'] = ['#include <ATen/cuda/CUDAHalf.cuh>']
env['extra_cuda_headers'].append('#include <ATen/DeviceGuard.h>')
sname = '' if scalar_name == "Float" else scalar_name
env['THType'] = 'Cuda{}'.format(sname)
env['THStorage'] = 'THCuda{}Storage'.format(sname)
env['THTensor'] = 'THCuda{}Tensor'.format(sname)
env['THIndexTensor'] = 'THCudaLongTensor'
env['state'] = ['globalContext().getTHCState()']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
env['Generator'] = 'CUDAGenerator'
else:
env['th_headers'] = [
'#include <TH/TH.h>',
'#include <TH/THTensor.hpp>',
'#include <THNN/THNN.h>',
'#undef THNN_',
]
env['extra_cuda_headers'] = []
env['THType'] = scalar_name
env['THStorage'] = "TH{}Storage".format(scalar_name)
env['THTensor'] = 'TH{}Tensor'.format(scalar_name)
env['THIndexTensor'] = 'THLongTensor'
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['Generator'] = 'CPUGenerator'
env['AS_REAL'] = env['ScalarType']
if scalar_name == "Half":
env['SparseTensor'] = 'Tensor'
if backend == "CUDA":
env['to_th_type'] = 'HalfFix<__half,Half>'
env['to_at_type'] = 'HalfFix<Half,__half>'
env['AS_REAL'] = 'convert<half,double>'
env['THScalarType'] = 'half'
else:
env['to_th_type'] = 'HalfFix<THHalf,Half>'
env['to_at_type'] = 'HalfFix<Half,THHalf>'
elif scalar_name == 'Long':
env['to_th_type'] = 'long'
env['to_at_type'] = 'int64_t'
else:
env['to_th_type'] = ''
env['to_at_type'] = ''
declarations, definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
fm = file_manager
if env['DenseBackend'] == 'CUDA':
fm = cuda_file_manager
if density != 'Sparse':
# there are no storage or tensor types for sparse; it's all uniform
fm.write(env['Storage'] + ".cpp", STORAGE_DERIVED_CPP, env)
fm.write(env['Storage'] + ".h", STORAGE_DERIVED_H, env)
env['TensorDenseOrSparse'] = TENSOR_DENSE_CPP.substitute(env)
fm.write(env['Tensor'] + ".cpp", TENSOR_DERIVED_CPP, env)
fm.write(env['Tensor'] + ".h", TENSOR_DERIVED_H, env)
if density != 'Sparse':
fm.write(env['Type'] + ".cpp", TYPE_DERIVED_CPP, env)
else:
fm.write(env['Type'] + ".cpp", SPARSE_TYPE_DERIVED_CPP, env)
fm.write(env['Type'] + ".h", TYPE_DERIVED_H, env)
type_register = TYPE_REGISTER.substitute(backend=env['Backend'],
scalar_type=scalar_name,
type_name=env['Type'])
if env['DenseBackend'] == 'CPU':
top_env['cpu_type_registrations'].append(type_register)
top_env['cpu_type_headers'].append('#include "ATen/{}.h"'.format(
env['Type']))
else:
assert env['DenseBackend'] == 'CUDA'
top_env['cuda_type_registrations'].append(type_register)
top_env['cuda_type_headers'].append('#include "ATen/{}.h"'.format(
env['Type']))
return env
def generate_storage_type_and_tensor(backend, scalar_type, declarations):
scalar_name, c_type, accreal, th_scalar_type = scalar_type
env = {}
env['ScalarName'] = scalar_name
env['ScalarType'] = c_type
env['THScalarType'] = th_scalar_type
env['AccScalarName'] = accreal
env['Storage'] = "{}{}Storage".format(backend, scalar_name)
env['Type'] = "{}{}Type".format(backend, scalar_name)
env['Tensor'] = "{}{}Tensor".format(backend, scalar_name)
env['Backend'] = backend
# used for generating switch logic for external functions
tag = backend + scalar_name
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
if backend == 'CUDA':
env['th_headers'] = [
'#include <THC/THC.h>', '#include <THCUNN/THCUNN.h>\n#undef THNN_'
]
sname = '' if scalar_name == "Float" else scalar_name
env['THType'] = 'Cuda{}'.format(sname)
env['THStorage'] = 'THCuda{}Storage'.format(sname)
env['THTensor'] = 'THCuda{}Tensor'.format(sname)
env['THIndexTensor'] = 'THCudaLongTensor'
env['state'] = ['context->thc_state']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
else:
env['th_headers'] = [
'#include <TH/TH.h>', '#include <THNN/THNN.h>\n#undef THNN_'
]
env['THType'] = scalar_name
env['THStorage'] = "TH{}Storage".format(scalar_name)
env['THTensor'] = 'TH{}Tensor'.format(scalar_name)
env['THIndexTensor'] = 'THLongTensor'
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['AS_REAL'] = env['ScalarType']
if scalar_name == "Half":
if backend == "CUDA":
env['to_th_half'] = 'HalfFix<__half,Half>'
env['to_at_half'] = 'HalfFix<Half,__half>'
env['AS_REAL'] = 'convert<half,double>'
env['THScalarType'] = 'half'
else:
env['to_th_half'] = 'HalfFix<THHalf,Half>'
env['to_at_half'] = 'HalfFix<Half,THHalf>'
else:
env['to_th_half'] = ''
env['to_at_half'] = ''
declarations, definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
write(env['Storage'] + ".cpp", STORAGE_DERIVED_CPP.substitute(env))
write(env['Storage'] + ".h", STORAGE_DERIVED_H.substitute(env))
write(env['Type'] + ".cpp", TYPE_DERIVED_CPP.substitute(env))
write(env['Type'] + ".h", TYPE_DERIVED_H.substitute(env))
write(env['Tensor'] + ".cpp", TENSOR_DERIVED_CPP.substitute(env))
write(env['Tensor'] + ".h", TENSOR_DERIVED_H.substitute(env))
type_register = (
('context->type_registry[static_cast<int>(Backend::{})]' +
'[static_cast<int>(ScalarType::{})].reset(new {}(context));').format(
backend, scalar_name, env['Type']))
top_env['type_registrations'].append(type_register)
top_env['type_headers'].append('#include "ATen/{}.h"'.format(env['Type']))
return env
def generate_storage_type_and_tensor(backend, density, declarations):
env = {}
density_tag = density if density != 'Dense' else ''
env['Density'] = density
env['Type'] = "{}{}Type".format(density_tag, backend)
env['DeviceType'] = backend_to_devicetype(backend)
env['Backend'] = density_tag + backend
env['storage_tensor_headers'] = []
if density != 'Sparse':
env['storage_tensor_headers'] = ['#include <c10/core/TensorImpl.h>']
# used for generating switch logic for external functions
tag = density_tag + backend
env['TypeID'] = 'TypeID::' + tag
top_env['type_ids'].append(tag + ',')
if backend == 'CUDA':
env['extra_cuda_headers'] = []
env['extra_cuda_headers'].append('#include <ATen/DeviceGuard.h>')
if options.rocm:
env['th_headers'] = [
'#include <THH/THH.h>',
'#include <THH/THHTensor.hpp>',
'#include <THHUNN/THHUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'].append('#include <ATen/hip/ATenHIPGeneral.h>')
env['extra_cuda_headers'].append('#include <ATen/hip/HIPDevice.h>')
env['extra_cuda_headers'].append('#include <ATen/hip/HIPTypeDefault.h>')
else:
env['th_headers'] = [
'#include <THC/THC.h>',
'#include <THC/THCTensor.hpp>',
'#include <THCUNN/THCUNN.h>',
'#undef THNN_',
'#undef THCIndexTensor_',
]
env['extra_cuda_headers'].append('#include <ATen/cuda/ATenCUDAGeneral.h>')
env['extra_cuda_headers'].append('#include <ATen/cuda/CUDADevice.h>')
env['extra_cuda_headers'].append('#include <ATen/cuda/CUDATypeDefault.h>')
env['state'] = ['globalContext().getTHCState()']
env['isCUDA'] = 'true'
env['storage_device'] = 'return storage->device;'
env['Generator'] = 'CUDAGenerator'
else:
env['th_headers'] = [
'#include <TH/TH.h>',
'#include <TH/THTensor.hpp>',
'#include <THNN/THNN.h>',
'#undef THNN_',
]
env['extra_cuda_headers'] = []
env['state'] = []
env['isCUDA'] = 'false'
env['storage_device'] = 'throw std::runtime_error("CPU storage has no device");'
env['Generator'] = 'CPUGenerator'
declarations, definitions = function_wrapper.create_derived(
env, declarations)
env['type_derived_method_declarations'] = declarations
env['type_derived_method_definitions'] = definitions
fm = file_manager
if env['DeviceType'] == 'CUDA':
fm = cuda_file_manager
if density != 'Sparse':
fm.write(env['Type'] + ".cpp", TYPE_DERIVED_CPP, env)
else:
fm.write(env['Type'] + ".cpp", SPARSE_TYPE_DERIVED_CPP, env)
fm.write(env['Type'] + ".h", TYPE_DERIVED_H, env)
type_register = TYPE_REGISTER.substitute(backend=env['Backend'], type_name=env['Type'])
if env['DeviceType'] == 'CPU':
top_env['cpu_type_registrations'].append(type_register)
top_env['cpu_type_headers'].append(
'#include "ATen/{}.h"'.format(env['Type']))
else:
assert env['DeviceType'] == 'CUDA'
top_env['cuda_type_registrations'].append(type_register)
top_env['cuda_type_headers'].append(
'#include "ATen/{}.h"'.format(env['Type']))
