| from string import Template |
| from copy import deepcopy |
| from . import CWrapPlugin |
| from itertools import product, chain |
| from collections import OrderedDict |
| |
| |
| class THPPlugin(CWrapPlugin): |
| |
| TYPE_UNPACK = { |
| 'THFloatTensor*': Template('((THPFloatTensor*)$arg)->cdata'), |
| 'THDoubleTensor*': Template('((THPDoubleTensor*)$arg)->cdata'), |
| 'THLongTensor*': Template('((THPLongTensor*)$arg)->cdata'), |
| 'THIntTensor*': Template('((THPIntTensor*)$arg)->cdata'), |
| 'THTensor*': Template('((THPTensor*)$arg)->cdata'), |
| 'THBoolTensor*': Template('((THPBoolTensor*)$arg)->cdata'), |
| 'THIndexTensor*': Template('((THPIndexTensor*)$arg)->cdata'), |
| 'THIntegerTensor*': Template('((THPIntegerTensor*)$arg)->cdata'), |
| |
| 'THCudaTensor*': Template('((THCPFloatTensor*)$arg)->cdata'), |
| 'THCudaDoubleTensor*': Template('((THCPDoubleTensor*)$arg)->cdata'), |
| 'THCudaIntTensor*': Template('((THCPIntTensor*)$arg)->cdata'), |
| 'THCudaLongTensor*': Template('((THCPLongTensor*)$arg)->cdata'), |
| |
| 'THSFloatTensor*': Template('((THSPFloatTensor*)$arg)->cdata'), |
| 'THSDoubleTensor*': Template('((THSPDoubleTensor*)$arg)->cdata'), |
| 'THSLongTensor*': Template('((THSPLongTensor*)$arg)->cdata'), |
| 'THSIntTensor*': Template('((THSPIntTensor*)$arg)->cdata'), |
| 'THSTensor*': Template('((THSPTensor*)$arg)->cdata'), |
| 'THSBoolTensor*': Template('((THSPBoolTensor*)$arg)->cdata'), |
| 'THSIndexTensor*': Template('((THSPIndexTensor*)$arg)->cdata'), |
| |
| 'THLongStorage*': Template('((THPLongStorage*)$arg)->cdata'), |
| 'THStorage*': Template('((THPStorage*)$arg)->cdata'), |
| 'THGenerator*': Template('THPGenerator_TH_CData($arg)'), |
| 'THSize*': Template('__size.get()'), |
| 'THStride*': Template('__stride.get()'), |
| 'void*': Template('THPUtils_unpackLong($arg)'), |
| 'long': Template('THPUtils_unpackLong($arg)'), |
| 'int': Template('((int) THPUtils_unpackLong($arg))'), |
| 'int64_t': Template('THPUtils_unpackLong($arg)'), |
| 'bool': Template('($arg == Py_True ? true : false)'), |
| 'float': Template('THPFloatUtils_unpackReal($arg)'), |
| 'double': Template('THPDoubleUtils_unpackReal($arg)'), |
| 'real': Template('THPUtils_(unpackReal)($arg)'), |
| 'accreal': Template('THPUtils_(unpackAccreal)($arg)'), |
| } |
| |
| TYPE_CHECK = { |
| 'THDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THPDoubleTensorClass'), |
| 'THFloatTensor*': Template('(PyObject*)Py_TYPE($arg) == THPFloatTensorClass'), |
| 'THLongTensor*': Template('(PyObject*)Py_TYPE($arg) == THPLongTensorClass'), |
| 'THIntTensor*': Template('(PyObject*)Py_TYPE($arg) == THPIntTensorClass'), |
| 'THTensor*': Template('(PyObject*)Py_TYPE($arg) == THPTensorClass'), |
| 'THBoolTensor*': Template('(PyObject*)Py_TYPE($arg) == THPBoolTensorClass'), |
| 'THIndexTensor*': Template('(PyObject*)Py_TYPE($arg) == THPIndexTensorClass'), |
| 'THIntegerTensor*': Template('(PyObject*)Py_TYPE($arg) == THPIntegerTensorClass'), |
| |
| 'THCudaTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPFloatTensorClass'), |
| 'THCudaDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPDoubleTensorClass'), |
| 'THCudaIntTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPIntTensorClass'), |
| 'THCudaLongTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPLongTensorClass'), |
| |
| 'THSDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPDoubleTensorClass'), |
| 'THSFloatTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPFloatTensorClass'), |
| 'THSLongTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPLongTensorClass'), |
| 'THSIntTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPIntTensorClass'), |
| 'THSTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPTensorClass'), |
| 'THSBoolTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPBoolTensorClass'), |
| 'THSIndexTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPIndexTensorClass'), |
| |
| 'THLongStorage*': Template('(PyObject*)Py_TYPE($arg) == THPLongStorageClass'), |
| 'THStorage*': Template('(PyObject*)Py_TYPE($arg) == THPStorageClass'), |
| 'THGenerator*': Template('(PyObject*)Py_TYPE($arg) == THPGeneratorClass'), |
| 'THSize*': Template('THPUtils_tryUnpackLongs($arg, __size)'), |
| 'THStride*': Template('THPUtils_tryUnpackLongs($arg, __stride)'), |
| 'void*': Template('THPUtils_checkLong($arg)'), |
| 'long': Template('THPUtils_checkLong($arg)'), |
| 'int64_t': Template('THPUtils_checkLong($arg)'), |
| 'int': Template('THPUtils_checkLong($arg)'), |
| 'bool': Template('PyBool_Check($arg)'), |
| 'float': Template('THPFloatUtils_checkReal($arg)'), |
| 'double': Template('THPDoubleUtils_checkReal($arg)'), |
| 'real': Template('THPUtils_(checkReal)($arg)'), |
| 'accreal': Template('THPUtils_(checkAccreal)($arg)'), |
| } |
| |
| SIZE_VARARG_CHECK = Template('THPUtils_tryUnpackLongVarArgs(args, $idx, __size)') |
| |
| RETURN_WRAPPER = { |
| 'THTensor*': Template('return THPTensor_(New)($result);'), |
| 'THSTensor*': Template('return THSPTensor_(New)($result);'), |
| 'THIndexTensor*': Template('return THPIndexTensor_(New)($result);'), |
| 'THLongTensor*': Template('return THPLongTensor_New($result);'), |
| 'THLongStorage*': Template('return THPLongStorage_New($result);'), |
| 'THCudaIntTensor*': Template('return THCPIntTensor_New($result);'), |
| 'THCudaLongTensor*': Template('return THCPLongTensor_New($result);'), |
| # TODO: make it smarter - it should return python long if result doesn't fit into an int |
| 'long': Template('return PyInt_FromLong($result);'), |
| 'int64_t': Template('return PyInt_FromLong($result);'), |
| 'int': Template('return PyLong_FromLong($result);'), |
| 'accreal': Template('return THPUtils_(newAccreal)($result);'), |
| 'self': Template('Py_INCREF(self);\nreturn (PyObject*)self;'), |
| 'real': Template('return THPUtils_(newReal)($result);'), |
| } |
| |
| TENSOR_METHODS_DECLARATION = Template(""" |
| static PyMethodDef TH${sparse}PTensor_$stateless(methods)[] = { |
| $methods |
| {NULL} |
| }; |
| """) |
| |
| WRAPPER_TEMPLATE = Template("""\ |
| PyObject * $name(PyObject *self, PyObject *args, PyObject *kwargs) |
| { |
| HANDLE_TH_ERRORS |
| int __tuplecount = args ? (int) PyTuple_Size(args) : 0; |
| int __dictcount = kwargs ? (int) PyDict_Size(kwargs) : 0; |
| int __argcount = __tuplecount + __dictcount; |
| $variables |
| $init |
| |
| $options |
| } |
| |
| THPUtils_invalidArguments(args, kwargs, "$readable_name", $num_options, $expected_args); |
| return NULL; |
| END_HANDLE_TH_ERRORS |
| } |
| """) |
| |
| ALLOCATE_TMPL = Template("""\ |
| THP${type}TensorPtr _${name}_guard((THP${type}Tensor*) THP${type}Tensor_NewEmpty()); |
| if (!_${name}_guard.get()) return NULL; |
| THP${type}Tensor* $name = _${name}_guard.get(); |
| """) |
| |
| ALLOCATE_CUDA = Template("""\ |
| #if IS_CUDA |
| ${cuda} |
| #else |
| ${cpu} |
| #endif |
| """) |
| |
| def _allocate(typename, tmpl, cuda_tmpl=None, sparse=False): |
| code = tmpl.safe_substitute(type=typename) |
| if typename == '': |
| code = code.replace('NewEmpty', '(NewEmpty)') |
| if cuda_tmpl: |
| cuda_code = code.replace('THP', 'THCP') |
| code = cuda_tmpl.substitute(cuda=cuda_code, cpu=code) |
| if sparse: |
| code = code.replace('THP', 'THSP') |
| code = code.replace('THCP', 'THCSP') |
| return Template(code) |
| |
| ALLOCATE_TYPE = { |
| 'THTensor*': _allocate('', ALLOCATE_TMPL), |
| 'THLongTensor*': _allocate('Long', ALLOCATE_TMPL), |
| 'THIntTensor*': _allocate('Int', ALLOCATE_TMPL), |
| 'THBoolTensor*': _allocate('Byte', ALLOCATE_TMPL, ALLOCATE_CUDA), |
| 'THIndexTensor*': _allocate('Long', ALLOCATE_TMPL, ALLOCATE_CUDA), |
| 'THIntegerTensor*': _allocate('Int', ALLOCATE_TMPL, ALLOCATE_CUDA), |
| |
| 'THSTensor*': _allocate('', ALLOCATE_TMPL, sparse=True), |
| } |
| |
| TYPE_NAMES = { |
| 'THTensor*': '" THPTensorStr "', |
| 'THSTensor*': '" THSPTensorStr "', |
| 'THStorage*': '" THPStorageStr "', |
| 'THGenerator*': 'torch.Generator', |
| 'THLongStorage*': '" THPModuleStr "LongStorage', |
| 'THLongTensor*': '" THPModuleStr "LongTensor', |
| 'THIntTensor*': '" THPModuleStr "IntTensor', |
| 'THBoolTensor*': '" THPModuleStr "ByteTensor', |
| 'THIndexTensor*': '" THPModuleStr "LongTensor', |
| 'THIntegerTensor*': '" THPModuleStr "IntTensor', |
| 'THFloatTensor*': '" THPModuleStr "FloatTensor', |
| 'THDoubleTensor*': '" THPModuleStr "DoubleTensor', |
| 'THCudaTensor*': 'torch.cuda.FloatTensor', |
| 'THCudaDoubleTensor*': 'torch.cuda.DoubleTensor', |
| 'THCudaIntTensor*': 'torch.cuda.IntTensor', |
| 'THCudaLongTensor*': 'torch.cuda.LongTensor', |
| 'THSize*': 'torch.Size', |
| 'THStride*': 'tuple', |
| 'long': 'int', |
| 'int64_t': 'int', |
| 'int': 'int', |
| 'real': '" RealStr "', |
| 'double': 'float', |
| 'accreal': '" RealStr "', |
| 'bool': 'bool', |
| 'const char*': 'bool', # Can come only from bool option. |
| } |
| |
| OUT_INIT = """ |
| ___out = kwargs ? PyDict_GetItemString(kwargs, "out") : NULL; |
| if (___out == Py_None) { ___out = NULL; __dictcount--; __argcount--; } |
| """ |
| |
| def __init__(self): |
| self.declarations = [] |
| self.stateless_declarations = [] |
| self.docstrings = [] |
| |
| BACKEND_SUBSTITUTIONS = { |
| 'CPU': 'TH', |
| 'CUDA': 'THCuda', |
| } |
| |
| def substitute_tensor_backend(self, arg, option): |
| if 'Backend' in arg['type']: |
| arg['type'] = arg['type'].replace('Backend', |
| self.BACKEND_SUBSTITUTIONS.get(option['backends'][0])) |
| # handle the fact that THCudaTensor isn't THCudaFloatTensor |
| if option['backends'][0] == 'CUDA' and 'Float' in arg['type']: |
| arg['type'] = arg['type'].replace('Float', '') |
| |
| def get_type_unpack(self, arg, option): |
| self.substitute_tensor_backend(arg, option) |
| return self.TYPE_UNPACK.get(arg['type'], None) |
| |
| def get_type_check(self, arg, option): |
| if arg['type'] == 'THSize*' and arg.get('long_args', False): |
| return self.SIZE_VARARG_CHECK |
| self.substitute_tensor_backend(arg, option) |
| return self.TYPE_CHECK.get(arg['type'], None) |
| |
| # TODO: argument descriptions shouldn't be part of THP, but rather a general cwrap thing |
| def get_wrapper_template(self, declaration): |
| arg_desc = OrderedDict() |
| |
| def format_arg(arg, var_args=False): |
| if var_args and arg.get('long_args', False): |
| return 'int ... ' + arg['name'] |
| else: |
| return self.TYPE_NAMES[arg['type']] + ' ' + arg['name'] |
| |
| def format_args(args, var_args=False): |
| option_desc = [format_arg(arg, var_args) |
| for arg in args |
| if not arg.get('ignore_check', False) and |
| not arg.get('output')] |
| output_args = list(filter(lambda a: a.get('output'), args)) |
| if output_args: |
| if len(output_args) > 1: |
| out_type = 'tuple[' |
| out_type += ', '.join( |
| self.TYPE_NAMES[arg['type']] for arg in output_args) |
| out_type += ']' |
| option_desc += ['#' + out_type + ' out'] |
| else: |
| arg = output_args[0] |
| option_desc += ['#' + self.TYPE_NAMES[arg['type']] + ' out'] |
| |
| if option_desc: |
| return '({})'.format(', '.join(option_desc)) |
| else: |
| return 'no arguments' |
| |
| for option in declaration['options']: |
| arg_desc[format_args(option['arguments'], False)] = True |
| arg_desc[format_args(option['arguments'], True)] = True |
| |
| arg_desc = sorted(list(arg_desc.keys()), key=len) |
| arg_desc = ['"' + desc + '"' for desc in arg_desc] |
| arg_str = ', '.join(arg_desc) |
| variables_str = '\n'.join(declaration.get('variables', [])) |
| init_str = '\n'.join(declaration.get('init', [])) |
| if 'stateless' in declaration['name']: |
| readable_name = 'torch.' + declaration['python_name'] |
| else: |
| readable_name = declaration['python_name'] |
| return Template(self.WRAPPER_TEMPLATE.safe_substitute( |
| readable_name=readable_name, num_options=len(arg_desc), |
| expected_args=arg_str, variables=variables_str, init=init_str)) |
| |
| def get_return_wrapper(self, option): |
| return self.RETURN_WRAPPER.get(option['return'], None) |
| |
| def get_arg_accessor(self, arg, option): |
| if arg['name'] == 'self': |
| return 'self' |
| if arg.get('output'): |
| if not option['output_provided']: |
| return arg['name'] |
| if option['output_count'] == 1: |
| return '___out' |
| else: |
| return 'PyTuple_GET_ITEM(___out, {})'.format(arg['output_idx']) |
| |
| def process_docstrings(self): |
| for declaration in self.declarations: |
| docstr = declaration.get('docstring_method') |
| if docstr is None: |
| continue |
| declaration['docstring_content'] = docstr.replace('\n', '\\n') |
| declaration['docstring_var'] = 'docstr_' + declaration['python_name'] |
| for declaration in self.stateless_declarations: |
| docstr = declaration.get('docstring_stateless') |
| if docstr is None: |
| continue |
| declaration['docstring_content'] = docstr.replace('\n', '\\n') |
| declaration['docstring_var'] = 'stateless_docstr_' + declaration['python_name'] |
| |
| def generate_out_options(self, declaration): |
| new_options = [] |
| declaration.setdefault('init', []) |
| declaration['init'] += [self.OUT_INIT] |
| for option in declaration['options']: |
| out_idx = [] |
| for i, arg in enumerate(option['arguments']): |
| if arg.get('output'): |
| out_idx.append(i) |
| if not out_idx: |
| option['has_output'] = True |
| option['output_provided'] = False |
| new_options.append(option) |
| continue |
| for output_provided in (True, False): |
| option_copy = deepcopy(option) |
| option_copy['has_output'] = True |
| option_copy['output_provided'] = output_provided |
| option_copy['output_count'] = len(out_idx) |
| for i, idx in enumerate(out_idx): |
| arg = option_copy['arguments'][idx] |
| arg['output_idx'] = i |
| if not output_provided: |
| arg['ignore_check'] = True |
| else: |
| option_copy['argcount_offset'] = -len(out_idx) + 1 |
| arg['no_kwargs'] = True |
| arg['no_idx'] = True |
| new_options.append(option_copy) |
| declaration['options'] = new_options |
| |
| def process_declarations(self, declarations): |
| new_declarations = [] |
| |
| def has_arg_type(declaration, type_name): |
| return any(arg['type'] == type_name |
| for option in declaration['options'] |
| for arg in option['arguments']) |
| |
| def has_long_args(declaration): |
| return any(arg.get('long_args', False) |
| for option in declaration['options'] |
| for arg in option['arguments']) |
| |
| def has_output_args(declaration): |
| return any(arg.get('output') |
| for option in declaration['options'] |
| for arg in option['arguments']) |
| |
| def backends_types_to_defined_if_string(declaration): |
| # A declaration has two fields: 'backend', which stores a list of |
| # backends (currently 'cpu' and 'cuda') the declaration applies |
| # to, and 'types', which stores a list of real types the |
| # declaration applies to. In PyTorch, when a function is only |
| # supported by a subset of types, we wrap it in macro definition |
| # checks. |
| # |
| # Previously, we manually required the cwrap declaration to |
| # specify for which backend/type combinations a function was |
| # defined for. Now, we explicitly list the types and backends for |
| # a declaration, if it should only be supported for a specific |
| # subset of types, backends, or type-backend pairs. |
| |
| types = declaration.get('types', []) |
| backends = declaration['backends'] |
| all_backends = ['CPU', 'CUDA'] |
| |
| def get_defined_string(backend, real): |
| if backend == 'CUDA': |
| if real == 'all': |
| return "IS_CUDA" |
| else: |
| return 'CUDA_{0}'.format(real.upper()) |
| else: |
| if real == 'all': |
| return "!IS_CUDA" |
| else: |
| return 'defined(TH_REAL_IS_{0})'.format(real.upper()) |
| |
| def expand_composite_type(p, t): |
| if t == 'floating_point': |
| result = ['double', 'float'] |
| if p == 'CUDA': |
| result.append('half') |
| elif t == 'integral': |
| result = ['byte', 'char', 'short', 'int', 'long'] |
| else: |
| result = [t] |
| return result |
| |
| defineds = [] |
| |
| # The logic below does not handle corner cases well. We allow the |
| # declaration to have a field 'backend_type_pairs' that stores a |
| # dictionary from type --> backend representing allowed |
| # combinations. Let's use these first. |
| for pair in declaration.get('backend_type_pairs', []): |
| p, t = pair |
| defineds.extend([get_defined_string(p, et) for et in |
| expand_composite_type(p, t)]) |
| |
| # In the base case, types is empty and backends contains both |
| # 'CPU' and 'CUDA' --> this means we support all types, and our |
| # string should be empty, or simply the list of explict type |
| # backend pairs |
| if (len(types) == 0 and all([proc in backends for proc in |
| all_backends])): |
| return " || ".join(defineds) |
| |
| # Case 2: types is empty, but only one backend type is specified |
| if len(types) == 0 and len(backends) == 1: |
| defineds.append('IS_CUDA' if backends[0] == 'CUDA' else |
| "!IS_CUDA") |
| return " || ".join(defineds) |
| |
| # Else, we loop overall all of the backend, type pairs and add |
| # them |
| for p in backends: |
| for t in types: |
| defineds.extend([get_defined_string(p, et) for et in |
| expand_composite_type(p, t)]) |
| |
| return " || ".join(defineds) |
| |
| for declaration in declarations: |
| # Disable all methods for THHalfTensor, unless cpu_half is True |
| |
| dfstr = backends_types_to_defined_if_string(declaration) |
| if len(dfstr) > 0: |
| # for now, need to check for distributed defined if as well |
| if 'defined_if' in declaration: |
| declaration['defined_if'] += ' && (' + dfstr + ')' |
| else: |
| declaration['defined_if'] = dfstr |
| |
| if not declaration.get('cpu_half', False): |
| defined_if = '!defined(TH_REAL_IS_HALF)' |
| if 'defined_if' in declaration: |
| defined_if += ' && (' + declaration['defined_if'] + ')' |
| declaration['defined_if'] = defined_if |
| |
| if declaration.get('only_register', False): |
| continue |
| |
| declaration.setdefault('python_name', declaration['name']) |
| declaration.setdefault('variables', []) |
| if has_arg_type(declaration, 'THSize*'): |
| declaration['variables'] += ['THLongStoragePtr __size;'] |
| if has_arg_type(declaration, 'THStride*'): |
| declaration['variables'] += ['THLongStoragePtr __stride;'] |
| if has_output_args(declaration): |
| declaration['variables'] += ['PyObject *___out;'] |
| self.generate_out_options(declaration) |
| if has_long_args(declaration): |
| for option in declaration['options']: |
| for arg in option['arguments']: |
| if arg.get('long_args', False): |
| arg['no_kwargs'] = True |
| for option in declaration['options']: |
| option['cname'] = 'TH{}Tensor_({})'.format( |
| 'S' if option.get('sparse', False) else '', option['cname']) |
| if option.get('sparse', False): |
| defined_if = option.get('defined_if', '') |
| option['defined_if'] = '!IS_DISTRIBUTED' + (' && ' if defined_if else '') + defined_if |
| |
| variants = declaration.get('variants', ['method']) |
| if 'function' in variants: |
| stateless_declaration = self.make_stateless(declaration) |
| new_declarations.append(stateless_declaration) |
| self.stateless_declarations.append(stateless_declaration) |
| if 'method' not in variants: |
| continue |
| |
| self.declarations.append(declaration) |
| declaration['name'] = 'TH{}PTensor_({})'.format( |
| 'S' if declaration.get('sparse', False) else '', declaration['name']) |
| for option in declaration['options']: |
| for arg in option['arguments']: |
| if arg['name'] == 'self': |
| arg['ignore_check'] = True |
| |
| register_only = [d for d in declarations if d.get('only_register', False)] |
| declarations = [d for d in declarations |
| if (('method' in d.get('variants', ['method'])) and |
| (not d.get('only_register', False)))] |
| self.declarations.extend(filter(lambda x: 'method' in x.get('variants', |
| ['method']), register_only)) |
| self.stateless_declarations.extend(filter(lambda x: 'method' not in |
| x.get('variants', ['method']), |
| register_only)) |
| |
| self.process_docstrings() |
| |
| all_declarations = declarations + new_declarations |
| return all_declarations |
| |
| def make_stateless(self, declaration): |
| declaration = deepcopy(declaration) |
| declaration['name'] = 'TH{}PTensor_stateless_({})'.format( |
| 'S' if declaration.get('sparse', False) else '', declaration['name']) |
| for option in declaration['options']: |
| for arg in option['arguments']: |
| if arg['name'] == 'self': |
| arg['assign_name'] = 'self' |
| arg['name'] = 'source' |
| return declaration |
| |
| def declare_methods(self, stateless, sparse): |
| tensor_methods = '' |
| for declaration in (self.declarations if not stateless else self.stateless_declarations): |
| if declaration.get('sparse', False) != sparse: |
| continue |
| flags = 'METH_VARARGS' |
| flags += ' | ' + declaration.get('method_flags') if 'method_flags' in declaration else '' |
| if not declaration.get('only_register'): |
| flags += ' | METH_KEYWORDS' |
| if declaration.get('override_method_flags'): |
| flags = declaration['override_method_flags'] |
| entry = Template(' {"$python_name", (PyCFunction)$name, $flags, $docstring},\n').substitute( |
| python_name=declaration['python_name'], name=declaration['name'], flags=flags, |
| docstring=declaration.get('docstring_var', 'NULL') |
| ) |
| if 'defined_if' in declaration: |
| entry = self.preprocessor_guard(entry, declaration['defined_if']) |
| tensor_methods += entry |
| generated = self.TENSOR_METHODS_DECLARATION.substitute( |
| methods=tensor_methods, |
| stateless=('' if not stateless else 'stateless_'), |
| sparse=('' if not sparse else 'S'), |
| ) |
| if sparse: |
| generated = '#if !defined(TH_REAL_IS_HALF) && !IS_DISTRIBUTED\n' + generated + '\n#endif\n\n' |
| return generated |
| |
| def process_full_file(self, code): |
| # We have to find a place before all undefs |
| idx = code.find('// PUT DEFINITIONS IN HERE PLEASE') |
| return (code[:idx] + |
| self.declare_methods(False, False) + |
| self.declare_methods(True, False) + |
| self.declare_methods(False, True) + |
| self.declare_methods(True, True) + |
| code[idx:] |
| ) |
| |
| def preprocessor_guard(self, code, condition): |
| return '#if ' + condition + '\n' + code + '#endif\n' |
| |
| def process_wrapper(self, code, declaration): |
| if 'defined_if' in declaration: |
| return self.preprocessor_guard(code, declaration['defined_if']) |
| return code |
| |
| def process_all_call_arg(self, code, option): |
| return 'LIBRARY_STATE ' + code |
| |
| def process_all_checks(self, code, option): |
| if option.get('has_output'): |
| indent = " " * 10 |
| if option['output_provided']: |
| checks = "___out != NULL &&\n" + indent |
| if option['output_count'] > 1: |
| checks += "PyTuple_Check(___out) &&\n" + indent |
| length_check = "PyTuple_GET_SIZE(___out) == {} &&\n".format( |
| option['output_count']) |
| checks += length_check + indent |
| code = checks + code |
| else: |
| code = "___out == NULL &&\n" + indent + code |
| |
| if any(arg.get('long_args', False) for arg in option['arguments']): |
| code = code.replace('__argcount ==', '__argcount >=') |
| expected = str(int(option.get('output_provided', False)) + |
| sum(not arg.get('no_kwargs', False) and not arg.get('ignore_check', False) |
| for arg in option['arguments'])) |
| code = '__dictcount == ' + expected + ' &&\n ' + code |
| |
| return code |
| |
| def process_option_code(self, code, option): |
| if option.get('defined_if', ''): |
| defined_if = option['defined_if'] |
| placeholder = '' |
| # This means that it's a first option, so we need a dummy if, |
| # so the next option can be an else if. |
| if 'else if' not in code: |
| placeholder = '\n #else\n if (false) {' |
| return '#if ' + defined_if + '\n ' + code + placeholder + '\n #endif\n' |
| return code |
| |
| def process_pre_arg_assign(self, template, option): |
| new_args = [] |
| for arg in option['arguments']: |
| if not option.get('output_provided', True) and arg.get('output'): |
| new_args.append(self.ALLOCATE_TYPE[arg['type']].substitute(name=arg['name'])) |
| template = new_args + template |
| return template |
| |
| def generate_docstrings_cpp(self): |
| template = Template('char* $name = "$content";') |
| return '\n\n'.join( |
| template.substitute(name=decl['docstring_var'], content=decl['docstring_content']) |
| for decl in chain(self.declarations, self.stateless_declarations) |
| if 'docstring_var' in decl) |
| |
| def generate_docstrings_h(self): |
| template = Template('extern char* $name;') |
| return '\n\n'.join( |
| template.substitute(name=decl['docstring_var']) |
| for decl in chain(self.declarations, self.stateless_declarations) |
| if 'docstring_var' in decl) |
