def generic(self, args, kws): assert not kws if len(args) == 1: # 0-dim arrays return one result array ary = args[0] ndim = max(ary.ndim, 1) retty = types.UniTuple(types.Array(types.intp, 1, 'C'), ndim) return signature(retty, ary) elif len(args) == 3: cond, x, y = args retdty = from_dtype( np.promote_types(as_dtype(getattr(args[1], 'dtype', args[1])), as_dtype(getattr(args[2], 'dtype', args[2])))) if isinstance(cond, types.Array): # array where() if isinstance(x, types.Array) and isinstance(y, types.Array): if (cond.ndim == x.ndim == y.ndim): if x.layout == y.layout == cond.layout: retty = types.Array(retdty, x.ndim, x.layout) else: retty = types.Array(retdty, x.ndim, 'C') return signature(retty, *args) else: # x and y both scalar retty = types.Array(retdty, cond.ndim, cond.layout) return signature(retty, *args) else: # scalar where() if not isinstance(x, types.Array): retty = types.Array(retdty, 0, 'C') return signature(retty, *args)
def _parse_nested_sequence(context, typ): """ Parse a (possibly 0d) nested sequence type. A (ndim, dtype) tuple is returned. Note the sequence may still be heterogeneous, as long as it converts to the given dtype. """ if isinstance(typ, (types.Buffer, )): raise TypingError("%r not allowed in a homogeneous sequence" % typ) elif isinstance(typ, (types.Sequence, )): n, dtype = _parse_nested_sequence(context, typ.dtype) return n + 1, dtype elif isinstance(typ, (types.BaseTuple, )): if typ.count == 0: # Mimick Numpy's behaviour return 1, types.float64 n, dtype = _parse_nested_sequence(context, typ[0]) dtypes = [dtype] for i in range(1, typ.count): _n, dtype = _parse_nested_sequence(context, typ[i]) if _n != n: raise TypingError("type %r does not have a regular shape" % (typ, )) dtypes.append(dtype) dtype = context.unify_types(*dtypes) if dtype is None: raise TypingError("cannot convert %r to a homogeneous type" % typ) return n + 1, dtype else: # Scalar type => check it's valid as a Numpy array dtype as_dtype(typ) return 0, typ
def expand_mapped_nb(mapped_arr: tp.Array1d, col_arr: tp.Array1d, idx_arr: tp.Array1d, target_shape: tp.Shape, fill_value: float) -> tp.Array2d: """Set each element to a value by boolean mask.""" nb_enabled = not isinstance(mapped_arr, np.ndarray) if nb_enabled: mapped_arr_dtype = as_dtype(mapped_arr.dtype) fill_value_dtype = as_dtype(fill_value) else: mapped_arr_dtype = mapped_arr.dtype fill_value_dtype = np.array(fill_value).dtype dtype = np.promote_types(mapped_arr_dtype, fill_value_dtype) def _expand_mapped_nb(mapped_arr, col_arr, idx_arr, target_shape, fill_value): out = np.full(target_shape, fill_value, dtype=dtype) for r in range(mapped_arr.shape[0]): out[idx_arr[r], col_arr[r]] = mapped_arr[r] return out if not nb_enabled: return _expand_mapped_nb(mapped_arr, col_arr, idx_arr, target_shape, fill_value) return _expand_mapped_nb
def stack_expand_mapped_nb(mapped_arr: tp.Array1d, col_map: tp.ColMap, fill_value: float) -> tp.Array2d: """Expand mapped array by stacking without using index data.""" nb_enabled = not isinstance(mapped_arr, np.ndarray) if nb_enabled: mapped_arr_dtype = as_dtype(mapped_arr.dtype) fill_value_dtype = as_dtype(fill_value) else: mapped_arr_dtype = mapped_arr.dtype fill_value_dtype = np.array(fill_value).dtype dtype = np.promote_types(mapped_arr_dtype, fill_value_dtype) def _stack_expand_mapped_nb(mapped_arr, col_map, fill_value): col_idxs, col_lens = col_map col_start_idxs = np.cumsum(col_lens) - col_lens out = np.full((np.max(col_lens), col_lens.shape[0]), fill_value, dtype=dtype) for col in range(col_lens.shape[0]): col_len = col_lens[col] if col_len == 0: continue col_start_idx = col_start_idxs[col] ridxs = col_idxs[col_start_idx:col_start_idx + col_len] out[:col_len, col] = mapped_arr[ridxs] return out if not nb_enabled: return _stack_expand_mapped_nb(mapped_arr, col_map, fill_value) return _stack_expand_mapped_nb
def ga_or(a, b): if isinstance(a, MultiVectorType) and isinstance(b, MultiVectorType): if a.layout_type != b.layout_type: raise numba.TypingError( "MultiVector objects belong to different layouts") imt_func = a.layout_type.obj.imt_func def impl(a, b): return a.layout.MultiVector(imt_func(a.value, b.value)) return impl elif isinstance(a, types.abstract.Number) and isinstance( b, MultiVectorType): ret_type = np.result_type(_numpy_support.as_dtype(a), _numpy_support.as_dtype(b.value_type.dtype)) def impl(a, b): return b.layout.MultiVector(np.zeros_like(b.value, dtype=ret_type)) return impl elif isinstance(a, MultiVectorType) and isinstance(b, types.abstract.Number): ret_type = np.result_type(_numpy_support.as_dtype(a.value_type.dtype), _numpy_support.as_dtype(b)) def impl(a, b): return a.layout.MultiVector(np.zeros_like(a.value, dtype=ret_type)) return impl
def ga_sub(a, b): if isinstance(a, MultiVectorType) and isinstance(b, MultiVectorType): if a.layout_type != b.layout_type: raise numba.TypingError( "MultiVector objects belong to different layouts") def impl(a, b): return a.layout.MultiVector(a.value - b.value) return impl elif isinstance(a, types.abstract.Number) and isinstance( b, MultiVectorType): scalar_index = b.layout_type.obj._basis_blade_order.bitmap_to_index[0] ret_type = np.result_type(_numpy_support.as_dtype(a), _numpy_support.as_dtype(b.value_type.dtype)) def impl(a, b): op = -b.value.astype(ret_type) op[scalar_index] += a return b.layout.MultiVector(op) return impl elif isinstance(a, MultiVectorType) and isinstance(b, types.abstract.Number): scalar_index = a.layout_type.obj._basis_blade_order.bitmap_to_index[0] ret_type = np.result_type(_numpy_support.as_dtype(a.value_type.dtype), _numpy_support.as_dtype(b)) def impl(a, b): op = a.value.astype(ret_type) op[scalar_index] -= b return a.layout.MultiVector(op) return impl
def find_common_dtype_from_numpy_dtypes(array_types, scalar_types): """Used to find common numba dtype for a sequences of numba dtypes each representing some numpy dtype""" np_array_dtypes = [numpy_support.as_dtype(dtype) for dtype in array_types] np_scalar_dtypes = [numpy_support.as_dtype(dtype) for dtype in scalar_types] np_common_dtype = numpy.find_common_type(np_array_dtypes, np_scalar_dtypes) numba_common_dtype = numpy_support.from_dtype(np_common_dtype) return numba_common_dtype
def kernel_wrapper(values): n = len(values) inputs = [np.empty(n, dtype=numpy_support.as_dtype(tp)) for tp in argtypes] output = np.empty(n, dtype=numpy_support.as_dtype(restype)) for i, vs in enumerate(values): for v, inp in zip(vs, inputs): inp[i] = v args = [output] + inputs kernel[int(math.ceil(n / 256)), 256](*args) return list(output)
def unify(self, typingctx, other): """ Unify the two number types using Numpy's rules. """ from numba.np import numpy_support if isinstance(other, Number): # XXX: this can produce unsafe conversions, # e.g. would unify {int64, uint64} to float64 a = numpy_support.as_dtype(self) b = numpy_support.as_dtype(other) sel = np.promote_types(a, b) return numpy_support.from_dtype(sel)
def check_round(cfunc, values, inty, outty, decimals): # Create input and output arrays of the right type arr = values.astype(as_dtype(inty)) out = np.zeros_like(arr).astype(as_dtype(outty)) pyout = out.copy() _fixed_np_round(arr, decimals, pyout) self.memory_leak_setup() cfunc(arr, decimals, out) self.memory_leak_teardown() np.testing.assert_allclose(out, pyout) # Output shape mismatch with self.assertRaises(ValueError) as raises: cfunc(arr, decimals, out[1:]) self.assertEqual(str(raises.exception), "invalid output shape")
def test_hypot(self, flags=enable_pyobj_flags): pyfunc = hypot x_types = [types.int64, types.uint64, types.float32, types.float64] x_values = [1, 2, 3, 4, 5, 6, .21, .34] y_values = [x + 2 for x in x_values] # Issue #563: precision issues with math.hypot() under Windows. prec = 'single' self.run_binary(pyfunc, x_types, x_values, y_values, flags, prec) # Check that values that overflow in naive implementations do not # in the numba impl def naive_hypot(x, y): return math.sqrt(x * x + y * y) for fltty in (types.float32, types.float64): cr = self.ccache.compile(pyfunc, (fltty, fltty), flags=flags) cfunc = cr.entry_point dt = numpy_support.as_dtype(fltty).type val = dt(np.finfo(dt).max / 30.) nb_ans = cfunc(val, val) self.assertPreciseEqual(nb_ans, pyfunc(val, val), prec='single') self.assertTrue(np.isfinite(nb_ans)) with warnings.catch_warnings(): warnings.simplefilter("error", RuntimeWarning) self.assertRaisesRegexp(RuntimeWarning, 'overflow encountered in .*_scalars', naive_hypot, val, val)
def test_generic_ptx(dtype): size = 500 lhs_arr = np.random.random(size).astype(dtype) lhs_col = Series(lhs_arr)._column rhs_arr = np.random.random(size).astype(dtype) rhs_col = Series(rhs_arr)._column def generic_function(a, b): return a ** 3 + b nb_type = numpy_support.from_dtype(cudf.dtype(dtype)) type_signature = (nb_type, nb_type) ptx_code, output_type = compile_ptx( generic_function, type_signature, device=True ) dtype = numpy_support.as_dtype(output_type).type out_col = libcudf.binaryop.binaryop_udf(lhs_col, rhs_col, ptx_code, dtype) result = lhs_arr ** 3 + rhs_arr np.testing.assert_almost_equal(result, out_col.to_array())
def build(self, cres): """ Returns (dtype numbers, function ptr, EnvironmentObject) """ # Buider wrapper for ufunc entry point signature = cres.signature info = build_gufunc_wrapper( self.py_func, cres, self.sin, self.sout, cache=self.cache, is_parfors=False, ) env = info.env ptr = info.library.get_pointer_to_function(info.name) # Get dtypes dtypenums = [] for a in signature.args: if isinstance(a, types.Array): ty = a.dtype else: ty = a dtypenums.append(as_dtype(ty).num) return dtypenums, ptr, env
def map_struct_to_record_dtype(cffi_type): """Convert a cffi type into a NumPy Record dtype """ fields = { 'names': [], 'formats': [], 'offsets': [], 'itemsize': ffi.sizeof(cffi_type), } is_aligned = True for k, v in cffi_type.fields: # guard unsupported values if v.bitshift != -1: msg = "field {!r} has bitshift, this is not supported" raise ValueError(msg.format(k)) if v.flags != 0: msg = "field {!r} has flags, this is not supported" raise ValueError(msg.format(k)) if v.bitsize != -1: msg = "field {!r} has bitsize, this is not supported" raise ValueError(msg.format(k)) dtype = numpy_support.as_dtype(map_type(v.type, use_record_dtype=True), ) fields['names'].append(k) fields['formats'].append(dtype) fields['offsets'].append(v.offset) # Check alignment is_aligned &= (v.offset % dtype.alignment == 0) return numpy_support.from_dtype(np.dtype(fields, align=is_aligned))
def compile_udf(udf, type_signature): """Compile ``udf`` with `numba` Compile a python callable function ``udf`` with `numba.cuda.compile_ptx_for_current_device(device=True)` using ``type_signature`` into CUDA PTX together with the generated output type. The output is expected to be passed to the PTX parser in `libcudf` to generate a CUDA device function to be inlined into CUDA kernels, compiled at runtime and launched. Parameters -------- udf: a python callable function type_signature: a tuple that specifies types of each of the input parameters of ``udf``. The types should be one in `numba.types` and could be converted from numpy types with `numba.numpy_support.from_dtype(...)`. Returns -------- ptx_code: The compiled CUDA PTX output_type: An numpy type """ # Check if we've already compiled a similar (but possibly distinct) # function before codebytes = udf.__code__.co_code if udf.__closure__ is not None: cvars = tuple([x.cell_contents for x in udf.__closure__]) cvarbytes = dumps(cvars) else: cvarbytes = b"" key = (type_signature, codebytes, cvarbytes) res = _udf_code_cache.get(key) if res: return res # We haven't compiled a function like this before, so need to fall back to # compilation with Numba ptx_code, return_type = cuda.compile_ptx_for_current_device(udf, type_signature, device=True) if not isinstance(return_type, cudf.core.udf.typing.MaskedType): output_type = numpy_support.as_dtype(return_type).type else: output_type = return_type # Populate the cache for this function res = (ptx_code, output_type) _udf_code_cache[key] = res return res
def get_default_scalar(self, nbtype): meta_type = type(nbtype) if isinstance(nbtype, types.Number): res = as_dtype(nbtype).type(self._default_data[meta_type]) elif isinstance(nbtype, types.UnicodeType): res = self._default_data[meta_type] return res
def test_generic_ptx(dtype): size = 500 lhs_arr = np.random.random(size).astype(dtype) lhs_col = Series(lhs_arr)._column rhs_arr = np.random.random(size).astype(dtype) rhs_col = Series(rhs_arr)._column @numba.cuda.jit(device=True) def generic_function(a, b): return a ** 3 + b nb_type = numpy_support.from_dtype(np.dtype(dtype)) type_signature = (nb_type, nb_type) result = generic_function.compile(type_signature) ptx = generic_function.inspect_ptx(type_signature) ptx_code = ptx.decode("utf-8") output_type = numpy_support.as_dtype(result.signature.return_type) out_col = libcudf.binaryop.binaryop_udf( lhs_col, rhs_col, ptx_code, output_type.type ) result = lhs_arr ** 3 + rhs_arr np.testing.assert_almost_equal(result, out_col.to_array())
def _get_proper_func(func_32, func_64, dtype, dist_name="the given"): """ Most of the standard NumPy distributions that accept dtype argument only support either np.float32 or np.float64 as dtypes. This is a helper function that helps Numba select the proper underlying implementation according to provided dtype. """ if isinstance(dtype, types.Omitted): dtype = dtype.value np_dt = dtype if isinstance(dtype, type): nb_dt = from_dtype(np.dtype(dtype)) elif isinstance(dtype, types.NumberClass): nb_dt = dtype np_dt = as_dtype(nb_dt) if np_dt not in [np.float32, np.float64]: raise TypingError("Argument dtype is not one of the" + " expected type(s): " + " np.float32 or np.float64") if np_dt == np.float32: next_func = func_32 else: next_func = func_64 return next_func, nb_dt
def compile_udf(udf, type_signature): """Compile ``udf`` with `numba` Compile a python callable function ``udf`` with `numba.cuda.compile_ptx_for_current_device(device=True)` using ``type_signature`` into CUDA PTX together with the generated output type. The output is expected to be passed to the PTX parser in `libcudf` to generate a CUDA device function to be inlined into CUDA kernels, compiled at runtime and launched. Parameters -------- udf: a python callable function type_signature: a tuple that specifies types of each of the input parameters of ``udf``. The types should be one in `numba.types` and could be converted from numpy types with `numba.numpy_support.from_dtype(...)`. Returns -------- ptx_code: The compiled CUDA PTX output_type: An numpy type """ ptx_code, return_type = cuda.compile_ptx_for_current_device(udf, type_signature, device=True) output_type = numpy_support.as_dtype(return_type) return (ptx_code, output_type.type)
def _build_element_wise_ufunc_wrapper(cres, signature): '''Build a wrapper for the ufunc loop entry point given by the compilation result object, using the element-wise signature. ''' ctx = cres.target_context library = cres.library fname = cres.fndesc.llvm_func_name with global_compiler_lock: info = build_ufunc_wrapper(library, ctx, fname, signature, cres.objectmode, cres) ptr = info.library.get_pointer_to_function(info.name) # Get dtypes dtypenums = [as_dtype(a).num for a in signature.args] dtypenums.append(as_dtype(signature.return_type).num) return dtypenums, ptr, cres.environment
def box_array(typ, val, c): assert c.context.enable_nrt np_dtype = numpy_support.as_dtype(typ.dtype) dtypeptr = c.env_manager.read_const(c.env_manager.add_const(np_dtype)) newary = c.pyapi.nrt_adapt_ndarray_to_python(typ, val, dtypeptr) # Steals NRT ref c.context.nrt.decref(c.builder, typ, val) return newary
def test_record_dtype_with_titles_roundtrip(self): recdtype = np.dtype([(("title a", 'a'), np.float), ('b', np.float)]) nbtype = numpy_support.from_dtype(recdtype) self.assertTrue(nbtype.is_title('title a')) self.assertFalse(nbtype.is_title('a')) self.assertFalse(nbtype.is_title('b')) got = numpy_support.as_dtype(nbtype) self.assertTrue(got, recdtype)
def get_random_sequence(self, nbtype, n=10): if isinstance(nbtype, types.Number): values = np.arange(n // 2, dtype=as_dtype(nbtype)) res = np.random.choice(values, n) elif isinstance(nbtype, types.UnicodeType): values = gen_strlist(n // 2) res = list(np.random.choice(values, n)) return res
def _get_py_col_dtype(ctype): """ Re-creates column dtype as python type to be used in read_csv call """ dtype = ctype.dtype if ctype == string_array_type: return str if isinstance(ctype, Categorical): return _reconstruct_CategoricalDtype(ctype.pd_dtype) return numpy_support.as_dtype(dtype)
def gen_array_factory(numba_dtype): """ Create a funtion that creates an array. Bind the numpy dtype because I don't know how to create the numba version """ np_dtype = numpy_support.as_dtype(numba_dtype) return lambda shape: np.zeros(shape, np_dtype)
def get_type_limits(eltype): np_dtype = numpy_support.as_dtype(eltype) if isinstance(eltype, types.Integer): return np.iinfo(np_dtype) elif isinstance(eltype, types.Float): return np.finfo(np_dtype) else: msg = 'Type {} not supported'.format(eltype) raise errors.TypingError(msg)
def compile_udf(udf, type_signature): """Compile ``udf`` with `numba` Compile a python callable function ``udf`` with `numba.cuda.compile_ptx_for_current_device(device=True)` using ``type_signature`` into CUDA PTX together with the generated output type. The output is expected to be passed to the PTX parser in `libcudf` to generate a CUDA device function to be inlined into CUDA kernels, compiled at runtime and launched. Parameters -------- udf: a python callable function type_signature: a tuple that specifies types of each of the input parameters of ``udf``. The types should be one in `numba.types` and could be converted from numpy types with `numba.numpy_support.from_dtype(...)`. Returns -------- ptx_code: The compiled CUDA PTX output_type: An numpy type """ import cudf.core.udf key = make_cache_key(udf, type_signature) res = _udf_code_cache.get(key) if res: return res # We haven't compiled a function like this before, so need to fall back to # compilation with Numba ptx_code, return_type = cuda.compile_ptx_for_current_device(udf, type_signature, device=True) if not isinstance(return_type, cudf.core.udf.typing.MaskedType): output_type = numpy_support.as_dtype(return_type).type else: output_type = return_type # Populate the cache for this function res = (ptx_code, output_type) _udf_code_cache[key] = res return res
def generic(self, args, kws): ufunc = self.ufunc base_types, explicit_outputs, ndims, layout = self._handle_inputs( ufunc, args, kws) ufunc_loop = ufunc_find_matching_loop(ufunc, base_types) if ufunc_loop is None: raise TypingError("can't resolve ufunc {0} for types {1}".format( ufunc.__name__, args)) # check if all the types involved in the ufunc loop are supported in this mode if not supported_ufunc_loop(ufunc, ufunc_loop): msg = "ufunc '{0}' using the loop '{1}' not supported in this mode" raise TypingError( msg=msg.format(ufunc.__name__, ufunc_loop.ufunc_sig)) # if there is any explicit output type, check that it is valid explicit_outputs_np = [as_dtype(tp.dtype) for tp in explicit_outputs] # Numpy will happily use unsafe conversions (although it will actually warn) if not all( np.can_cast(fromty, toty, 'unsafe') for (fromty, toty ) in zip(ufunc_loop.numpy_outputs, explicit_outputs_np)): msg = "ufunc '{0}' can't cast result to explicit result type" raise TypingError(msg=msg.format(ufunc.__name__)) # A valid loop was found that is compatible. The result of type inference should # be based on the explicit output types, and when not available with the type given # by the selected NumPy loop out = list(explicit_outputs) implicit_output_count = ufunc.nout - len(explicit_outputs) if implicit_output_count > 0: # XXX this is sometimes wrong for datetime64 and timedelta64, # as ufunc_find_matching_loop() doesn't do any type inference ret_tys = ufunc_loop.outputs[-implicit_output_count:] if ndims > 0: assert layout is not None ret_tys = [ types.Array(dtype=ret_ty, ndim=ndims, layout=layout) for ret_ty in ret_tys ] ret_tys = [ resolve_output_type(self.context, args, ret_ty) for ret_ty in ret_tys ] out.extend(ret_tys) # note: although the previous code should support multiple return values, only one # is supported as of now (signature may not support more than one). # there is an check enforcing only one output out.extend(args) return signature(*out)
def dtype_generated_usecase(a, b, dtype=None): if isinstance(dtype, (types.misc.NoneType, types.misc.Omitted)): out_dtype = np.result_type(*(np.dtype(ary.dtype.name) for ary in (a, b))) elif isinstance(dtype, (types.DType, types.NumberClass)): out_dtype = as_dtype(dtype) else: raise TypeError("Unhandled Type %s" % type(dtype)) def _fn(a, b, dtype=None): return np.ones(a.shape, dtype=out_dtype) return _fn
def box_array(typ, val, c): nativearycls = c.context.make_array(typ) nativeary = nativearycls(c.context, c.builder, value=val) if c.context.enable_nrt: np_dtype = numpy_support.as_dtype(typ.dtype) dtypeptr = c.env_manager.read_const(c.env_manager.add_const(np_dtype)) # Steals NRT ref newary = c.pyapi.nrt_adapt_ndarray_to_python(typ, val, dtypeptr) return newary else: parent = nativeary.parent c.pyapi.incref(parent) return parent