"""Common CUDA related functionality. """ from __future__ import print_function import logging from pytools import Record, RecordWithoutPickling from pytools.persistent_dict import KeyBuilder as KeyBuilderBase from pytools.persistent_dict import WriteOncePersistentDict from pycuda._cluda import CLUDA_PREAMBLE import pycuda._mymako as mako from pycuda.tools import (dtype_to_ctype, bitlog2, context_dependent_memoize, ScalarArg, VectorArg) import pycuda.gpuarray as gpuarray # from compyle.thrust.sort import argsort import pycuda.driver as drv from pycuda.compiler import SourceModule as _SourceModule from pytools import memoize import numpy as np import six _cuda_ctx = False def set_context(): global _cuda_ctx if not _cuda_ctx: import pycuda.autoinit _cuda_ctx = True # The following code is taken from pyopencl for struct mapping. # it should be ported over to pycuda eventually. import pycuda.gpuarray as gpuarray # noqa class SourceModule(_SourceModule): def __getattr__(self, name): def kernel(*args, **kwargs): f = self.get_function(name) return f(*args, **kwargs) kernel.function_name = name return kernel class _CDeclList: def __init__(self, device): self.device = device self.declared_dtypes = set() self.declarations = [] self.saw_complex = False def add_dtype(self, dtype): dtype = np.dtype(dtype) if dtype.kind == "c": self.saw_complex = True if dtype.kind != "V": return if dtype in self.declared_dtypes: return for name, field_data in sorted(six.iteritems(dtype.fields)): field_dtype, offset = field_data[:2] self.add_dtype(field_dtype) _, cdecl = match_dtype_to_c_struct( self.device, dtype_to_ctype(dtype), dtype ) self.declarations.append(cdecl) self.declared_dtypes.add(dtype) def visit_arguments(self, arguments): for arg in arguments: dtype = arg.dtype if dtype.kind == "c": self.saw_complex = True def get_declarations(self): result = "\n\n".join(self.declarations) if self.saw_complex: result = ( "#include \n\n" + result) return result @memoize def match_dtype_to_c_struct(device, name, dtype, context=None, use_typedef=False): """Return a tuple `(dtype, c_decl)` such that the C struct declaration in `c_decl` and the structure :class:`numpy.dtype` instance `dtype` have the same memory layout. Note that *dtype* may be modified from the value that was passed in, for example to insert padding. (As a remark on implementation, this routine runs a small kernel on the given *device* to ensure that :mod:`numpy` and C offsets and sizes match.) This example explains the use of this function:: >>> import numpy as np >>> import pyopencl as cl >>> import pyopencl.tools >>> ctx = cl.create_some_context() >>> dtype = np.dtype([("id", np.uint32), ("value", np.float32)]) >>> dtype, c_decl = pyopencl.tools.match_dtype_to_c_struct( ... ctx.devices[0], 'id_val', dtype) >>> print c_decl typedef struct { unsigned id; float value; } id_val; >>> print dtype [('id', '>> cl.tools.get_or_register_dtype('id_val', dtype) As this example shows, it is important to call :func:`get_or_register_dtype` on the modified `dtype` returned by this function, not the original one. """ fields = sorted( six.iteritems(dtype.fields), key=lambda name_dtype_offset: name_dtype_offset[1][1] ) c_fields = [] for field_name, dtype_and_offset in fields: field_dtype, offset = dtype_and_offset[:2] c_fields.append(" %s %s;" % (dtype_to_ctype(field_dtype), field_name)) if use_typedef: c_decl = "typedef struct {\n%s\n} %s;\n\n" % ( "\n".join(c_fields), name ) else: c_decl = "struct %s {\n%s\n};\n\n" % ( name, "\n".join(c_fields) ) cdl = _CDeclList(device) for field_name, dtype_and_offset in fields: field_dtype, offset = dtype_and_offset[:2] cdl.add_dtype(field_dtype) pre_decls = cdl.get_declarations() offset_code = "\n".join( "result[%d] = pycuda_offsetof(%s, %s);" % (i + 1, name, field_name) for i, (field_name, _) in enumerate(fields)) src = r""" #define pycuda_offsetof(st, m) \ ((uint) ((char *) &(dummy_pycuda.m) \ - (char *)&dummy_pycuda )) %(pre_decls)s %(my_decl)s extern "C" __global__ void get_size_and_offsets(uint *result) { result[0] = sizeof(%(my_type)s); %(my_type)s dummy_pycuda; %(offset_code)s } """ % dict( pre_decls=pre_decls, my_decl=c_decl, my_type=name, offset_code=offset_code) prg = SourceModule(src) knl = prg.get_size_and_offsets result_buf = gpuarray.empty(1 + len(fields), np.uint32) e = drv.Event() knl(result_buf.gpudata, block=(1, 1, 1)) e.record() e.synchronize() size_and_offsets = result_buf.get() size = int(size_and_offsets[0]) from pytools import any offsets = size_and_offsets[1:] if any(ofs >= size for ofs in offsets): # offsets not plausible if dtype.itemsize == size: # If sizes match, use numpy's idea of the offsets. offsets = [dtype_and_offset[1] for field_name, dtype_and_offset in fields] else: raise RuntimeError( "OpenCL compiler reported offsetof() past sizeof() " "for struct layout on '%s'. " "This makes no sense, and it's usually indicates a " "compiler bug. " "Refusing to discover struct layout." % device) del knl del prg del context try: dtype_arg_dict = { 'names': [field_name for field_name, (field_dtype, offset) in fields], 'formats': [field_dtype for field_name, (field_dtype, offset) in fields], 'offsets': [int(x) for x in offsets], 'itemsize': int(size_and_offsets[0]), } dtype = np.dtype(dtype_arg_dict) if dtype.itemsize != size_and_offsets[0]: # "Old" versions of numpy (1.6.x?) silently ignore "itemsize". Boo. dtype_arg_dict["names"].append("_pycl_size_fixer") dtype_arg_dict["formats"].append(np.uint8) dtype_arg_dict["offsets"].append(int(size_and_offsets[0]) - 1) dtype = np.dtype(dtype_arg_dict) except NotImplementedError: def calc_field_type(): total_size = 0 padding_count = 0 for offset, (field_name, (field_dtype, _)) in zip(offsets, fields): if offset > total_size: padding_count += 1 yield ('__pycuda_padding%d' % padding_count, 'V%d' % offset - total_size) yield field_name, field_dtype total_size = field_dtype.itemsize + offset dtype = np.dtype(list(calc_field_type())) assert dtype.itemsize == size_and_offsets[0] return dtype, c_decl @memoize def dtype_to_c_struct(device, dtype): if dtype.fields is None: return "" import pyopencl.cltypes if dtype in pyopencl.cltypes.vec_type_to_scalar_and_count: # Vector types are built-in. Don't try to redeclare those. return "" matched_dtype, c_decl = match_dtype_to_c_struct( device, dtype_to_ctype(dtype), dtype) def dtypes_match(): result = len(dtype.fields) == len(matched_dtype.fields) for name, val in six.iteritems(dtype.fields): result = result and matched_dtype.fields[name] == val return result assert dtypes_match() return c_decl ##################################################################### # The GenericScanKernel is added here temporarily until the following # PR is merged into PyCUDA # https://github.com/inducer/pycuda/pull/188 ##################################################################### logger = logging.getLogger(__name__) ##################################################################### # The GenericScanKernel is added here temporarily until the following # PR is merged into PyCUDA # https://github.com/inducer/pycuda/pull/188 ##################################################################### def parse_arg_list(arguments): """Parse a list of kernel arguments. *arguments* may be a comma-separate list of C declarators in a string, a list of strings representing C declarators, or :class:`Argument` objects. """ if isinstance(arguments, str): arguments = arguments.split(",") def parse_single_arg(obj): if isinstance(obj, str): from pycuda.tools import parse_c_arg return parse_c_arg(obj) else: return obj return [parse_single_arg(arg) for arg in arguments] def get_arg_list_scalar_arg_dtypes(arg_types): result = [] for arg_type in arg_types: if isinstance(arg_type, ScalarArg): result.append(arg_type.dtype) elif isinstance(arg_type, VectorArg): result.append(None) else: raise RuntimeError("arg type not understood: %s" % type(arg_type)) return result def _process_code_for_macro(code): if "//" in code: raise RuntimeError("end-of-line comments ('//') may not be used in " "code snippets") return code.replace("\n", " \\\n") class _NumpyTypesKeyBuilder(KeyBuilderBase): def update_for_type(self, key_hash, key): if issubclass(key, np.generic): self.update_for_str(key_hash, key.__name__) return raise TypeError("unsupported type for persistent hash keying: %s" % type(key)) # {{{ preamble SHARED_PREAMBLE = CLUDA_PREAMBLE + """ #define WG_SIZE ${wg_size} #define SCAN_EXPR(a, b, across_seg_boundary) ${scan_expr} #define INPUT_EXPR(i) (${input_expr}) %if is_segmented: #define IS_SEG_START(i, a) (${is_segment_start_expr}) %endif ${preamble} typedef ${dtype_to_ctype(scan_dtype)} scan_type; typedef ${dtype_to_ctype(index_dtype)} index_type; // NO_SEG_BOUNDARY is the largest representable integer in index_type. // This assumption is used in code below. #define NO_SEG_BOUNDARY ${str(np.iinfo(index_dtype).max)} """ # }}} # {{{ main scan code # Algorithm: Each work group is responsible for one contiguous # 'interval'. There are just enough intervals to fill all compute # units. Intervals are split into 'units'. A unit is what gets # worked on in parallel by one work group. # # in index space: # interval > unit > local-parallel > k-group # # (Note that there is also a transpose in here: The data is read # with local ids along linear index order.) # # Each unit has two axes--the local-id axis and the k axis. # # unit 0: # | | | | | | | | | | ----> lid # | | | | | | | | | | # | | | | | | | | | | # | | | | | | | | | | # | | | | | | | | | | # # | # v k (fastest-moving in linear index) # # unit 1: # | | | | | | | | | | ----> lid # | | | | | | | | | | # | | | | | | | | | | # | | | | | | | | | | # | | | | | | | | | | # # | # v k (fastest-moving in linear index) # # ... # # At a device-global level, this is a three-phase algorithm, in # which first each interval does its local scan, then a scan # across intervals exchanges data globally, and the final update # adds the exchanged sums to each interval. # # Exclusive scan is realized by allowing look-behind (access to the # preceding item) in the final update, by means of a local shift. # # NOTE: All segment_start_in_X indices are relative to the start # of the array. SCAN_INTERVALS_SOURCE = SHARED_PREAMBLE + r""" #define K ${k_group_size} // #define DEBUG #ifdef DEBUG #define pycu_printf(ARGS) printf ARGS #else #define pycu_printf(ARGS) /* */ #endif KERNEL REQD_WG_SIZE(WG_SIZE, 1, 1) void ${kernel_name}( ${argument_signature}, GLOBAL_MEM scan_type* __restrict__ partial_scan_buffer, const index_type N, const index_type interval_size %if is_first_level: , GLOBAL_MEM scan_type* __restrict__ interval_results %endif %if is_segmented and is_first_level: // NO_SEG_BOUNDARY if no segment boundary in interval. , GLOBAL_MEM index_type* __restrict__ g_first_segment_start_in_interval %endif %if store_segment_start_flags: , GLOBAL_MEM char* __restrict__ g_segment_start_flags %endif ) { // index K in first dimension used for carry storage %if use_bank_conflict_avoidance: // Avoid bank conflicts by adding a single 32-bit value to the size of // the scan type. struct __attribute__ ((__packed__)) wrapped_scan_type { scan_type value; int dummy; }; %else: struct wrapped_scan_type { scan_type value; }; %endif // padded in WG_SIZE to avoid bank conflicts LOCAL_MEM struct wrapped_scan_type ldata[K + 1][WG_SIZE + 1]; %if is_segmented: LOCAL_MEM char l_segment_start_flags[K][WG_SIZE]; LOCAL_MEM index_type l_first_segment_start_in_subtree[WG_SIZE]; // only relevant/populated for local id 0 index_type first_segment_start_in_interval = NO_SEG_BOUNDARY; index_type first_segment_start_in_k_group, first_segment_start_in_subtree; %endif // {{{ declare local data for input_fetch_exprs if any of them are stenciled <% fetch_expr_offsets = {} for name, arg_name, ife_offset in input_fetch_exprs: fetch_expr_offsets.setdefault(arg_name, set()).add(ife_offset) local_fetch_expr_args = set( arg_name for arg_name, ife_offsets in fetch_expr_offsets.items() if -1 in ife_offsets or len(ife_offsets) > 1) %> %for arg_name in local_fetch_expr_args: LOCAL_MEM ${arg_ctypes[arg_name]} l_${arg_name}[WG_SIZE*K]; %endfor // }}} const index_type interval_begin = interval_size * GID_0; const index_type interval_end = min(interval_begin + interval_size, N); const index_type unit_size = K * WG_SIZE; index_type unit_base = interval_begin; %for is_tail in [False, True]: %if not is_tail: for(; unit_base + unit_size <= interval_end; unit_base += unit_size) %else: if (unit_base < interval_end) %endif { // {{{ carry out input_fetch_exprs // (if there are ones that need to be fetched into local) %if local_fetch_expr_args: for(index_type k = 0; k < K; k++) { const index_type offset = k*WG_SIZE + LID_0; const index_type read_i = unit_base + offset; %for arg_name in local_fetch_expr_args: %if is_tail: if (read_i < interval_end) %endif { l_${arg_name}[offset] = ${arg_name}[read_i]; } %endfor } local_barrier(); %endif pycu_printf(("after input_fetch_exprs\n")); // }}} // {{{ read a unit's worth of data from global for(index_type k = 0; k < K; k++) { const index_type offset = k*WG_SIZE + LID_0; const index_type read_i = unit_base + offset; %if is_tail: if (read_i < interval_end) %endif { %for name, arg_name, ife_offset in input_fetch_exprs: ${arg_ctypes[arg_name]} ${name}; %if arg_name in local_fetch_expr_args: if (offset + ${ife_offset} >= 0) ${name} = l_${arg_name}[offset + ${ife_offset}]; else if (read_i + ${ife_offset} >= 0) ${name} = ${arg_name}[read_i + ${ife_offset}]; /* else if out of bounds, name is left undefined */ %else: // ${arg_name} gets fetched directly from global ${name} = ${arg_name}[read_i]; %endif %endfor scan_type scan_value = INPUT_EXPR(read_i); const index_type o_mod_k = offset % K; const index_type o_div_k = offset / K; ldata[o_mod_k][o_div_k].value = scan_value; %if is_segmented: bool is_seg_start = IS_SEG_START(read_i, scan_value); l_segment_start_flags[o_mod_k][o_div_k] = is_seg_start; %endif %if store_segment_start_flags: g_segment_start_flags[read_i] = is_seg_start; %endif } } pycu_printf(("after read from global\n")); // }}} // {{{ carry in from previous unit, if applicable %if is_segmented: local_barrier(); first_segment_start_in_k_group = NO_SEG_BOUNDARY; if (l_segment_start_flags[0][LID_0]) first_segment_start_in_k_group = unit_base + K*LID_0; %endif if (LID_0 == 0 && unit_base != interval_begin) { scan_type tmp = ldata[K][WG_SIZE - 1].value; scan_type tmp_aux = ldata[0][0].value; ldata[0][0].value = SCAN_EXPR( tmp, tmp_aux, %if is_segmented: (l_segment_start_flags[0][0]) %else: false %endif ); } pycu_printf(("after carry-in\n")); // }}} local_barrier(); // {{{ scan along k (sequentially in each work item) scan_type sum = ldata[0][LID_0].value; %if is_tail: const index_type offset_end = interval_end - unit_base; %endif for(index_type k = 1; k < K; k++) { %if is_tail: if (K * LID_0 + k < offset_end) %endif { scan_type tmp = ldata[k][LID_0].value; %if is_segmented: index_type seq_i = unit_base + K*LID_0 + k; if (l_segment_start_flags[k][LID_0]) { first_segment_start_in_k_group = min( first_segment_start_in_k_group, seq_i); } %endif sum = SCAN_EXPR(sum, tmp, %if is_segmented: (l_segment_start_flags[k][LID_0]) %else: false %endif ); ldata[k][LID_0].value = sum; } } pycu_printf(("after scan along k\n")); // }}} // store carry in out-of-bounds (padding) array entry (index K) in // the K direction ldata[K][LID_0].value = sum; %if is_segmented: l_first_segment_start_in_subtree[LID_0] = first_segment_start_in_k_group; %endif local_barrier(); // {{{ tree-based local parallel scan // This tree-based scan works as follows: // - Each work item adds the previous item to its current state // - barrier // - Each work item adds in the item from two positions to the left // - barrier // - Each work item adds in the item from four positions to the left // ... // At the end, each item has summed all prior items. // across k groups, along local id // (uses out-of-bounds k=K array entry for storage) scan_type val = ldata[K][LID_0].value; <% scan_offset = 1 %> % while scan_offset <= wg_size: // {{{ reads from local allowed, writes to local not allowed if (LID_0 >= ${scan_offset}) { scan_type tmp = ldata[K][LID_0 - ${scan_offset}].value; % if is_tail: if (K*LID_0 < offset_end) % endif { val = SCAN_EXPR(tmp, val, %if is_segmented: (l_first_segment_start_in_subtree[LID_0] != NO_SEG_BOUNDARY) %else: false %endif ); } %if is_segmented: // Prepare for l_first_segment_start_in_subtree, below. // Note that this update must take place *even* if we're // out of bounds. first_segment_start_in_subtree = min( l_first_segment_start_in_subtree[LID_0], l_first_segment_start_in_subtree [LID_0 - ${scan_offset}]); %endif } %if is_segmented: else { first_segment_start_in_subtree = l_first_segment_start_in_subtree[LID_0]; } %endif // }}} local_barrier(); // {{{ writes to local allowed, reads from local not allowed ldata[K][LID_0].value = val; %if is_segmented: l_first_segment_start_in_subtree[LID_0] = first_segment_start_in_subtree; %endif // }}} local_barrier(); %if 0: if (LID_0 == 0) { printf("${scan_offset}: "); for (int i = 0; i < WG_SIZE; ++i) { if (l_first_segment_start_in_subtree[i] == NO_SEG_BOUNDARY) printf("- "); else printf("%d ", l_first_segment_start_in_subtree[i]); } printf("\n"); } %endif <% scan_offset *= 2 %> % endwhile pycu_printf(("after tree scan\n")); // }}} // {{{ update local values if (LID_0 > 0) { sum = ldata[K][LID_0 - 1].value; for(index_type k = 0; k < K; k++) { %if is_tail: if (K * LID_0 + k < offset_end) %endif { scan_type tmp = ldata[k][LID_0].value; ldata[k][LID_0].value = SCAN_EXPR(sum, tmp, %if is_segmented: (unit_base + K * LID_0 + k >= first_segment_start_in_k_group) %else: false %endif ); } } } %if is_segmented: if (LID_0 == 0) { // update interval-wide first-seg variable from current unit first_segment_start_in_interval = min( first_segment_start_in_interval, l_first_segment_start_in_subtree[WG_SIZE-1]); } %endif pycu_printf(("after local update\n")); // }}} local_barrier(); // {{{ write data { // work hard with index math to achieve contiguous 32-bit stores GLOBAL_MEM int *dest = (GLOBAL_MEM int *) (partial_scan_buffer + unit_base); <% assert scan_dtype.itemsize % 4 == 0 ints_per_wg = wg_size ints_to_store = scan_dtype.itemsize*wg_size*k_group_size // 4 %> const index_type scan_types_per_int = ${scan_dtype.itemsize//4}; %for store_base in range(0, ints_to_store, ints_per_wg): <% # Observe that ints_to_store is divisible by the work group # size already, so we won't go out of bounds that way. assert store_base + ints_per_wg <= ints_to_store %> %if is_tail: if (${store_base} + LID_0 < scan_types_per_int*(interval_end - unit_base)) %endif { index_type linear_index = ${store_base} + LID_0; index_type linear_scan_data_idx = linear_index / scan_types_per_int; index_type remainder = linear_index - linear_scan_data_idx * scan_types_per_int; int* src = (int*) &(ldata [linear_scan_data_idx % K] [linear_scan_data_idx / K].value); dest[linear_index] = src[remainder]; } %endfor } pycu_printf(("after write\n")); // }}} local_barrier(); } % endfor // write interval sum %if is_first_level: if (LID_0 == 0) { interval_results[GID_0] = partial_scan_buffer[interval_end - 1]; %if is_segmented: g_first_segment_start_in_interval[GID_0] = first_segment_start_in_interval; %endif } %endif } """ # }}} # {{{ update UPDATE_SOURCE = SHARED_PREAMBLE + r""" KERNEL REQD_WG_SIZE(WG_SIZE, 1, 1) void ${name_prefix}_final_update( ${argument_signature}, const index_type N, const index_type interval_size, GLOBAL_MEM scan_type* __restrict__ interval_results, GLOBAL_MEM scan_type* __restrict__ partial_scan_buffer %if is_segmented: , GLOBAL_MEM index_type* __restrict__ g_first_segment_start_in_interval %endif %if is_segmented and use_lookbehind_update: , GLOBAL_MEM char* __restrict__ g_segment_start_flags %endif ) { %if use_lookbehind_update: LOCAL_MEM scan_type ldata[WG_SIZE]; %endif %if is_segmented and use_lookbehind_update: LOCAL_MEM char l_segment_start_flags[WG_SIZE]; %endif const index_type interval_begin = interval_size * GID_0; const index_type interval_end = min(interval_begin + interval_size, N); // carry from last interval scan_type carry = ${neutral}; if (GID_0 != 0) carry = interval_results[GID_0 - 1]; %if is_segmented: const index_type first_seg_start_in_interval = g_first_segment_start_in_interval[GID_0]; %endif %if not is_segmented and 'last_item' in output_statement: scan_type last_item = interval_results[GDIM_0-1]; %endif %if not use_lookbehind_update: // {{{ no look-behind ('prev_item' not in output_statement -> simpler) index_type update_i = interval_begin+LID_0; %if is_segmented: index_type seg_end = min(first_seg_start_in_interval, interval_end); %endif for(; update_i < interval_end; update_i += WG_SIZE) { scan_type partial_val = partial_scan_buffer[update_i]; scan_type item = SCAN_EXPR(carry, partial_val, %if is_segmented: (update_i >= seg_end) %else: false %endif ); index_type i = update_i; { ${output_statement}; } } // }}} %else: // {{{ allow look-behind ('prev_item' in output_statement -> complicated) // We are not allowed to branch across barriers at a granularity smaller // than the whole workgroup. Therefore, the for loop is group-global, // and there are lots of local ifs. index_type group_base = interval_begin; scan_type prev_item = carry; // (A) for(; group_base < interval_end; group_base += WG_SIZE) { index_type update_i = group_base+LID_0; // load a work group's worth of data if (update_i < interval_end) { scan_type tmp = partial_scan_buffer[update_i]; tmp = SCAN_EXPR(carry, tmp, %if is_segmented: (update_i >= first_seg_start_in_interval) %else: false %endif ); ldata[LID_0] = tmp; %if is_segmented: l_segment_start_flags[LID_0] = g_segment_start_flags[update_i]; %endif } local_barrier(); // find prev_item if (LID_0 != 0) prev_item = ldata[LID_0 - 1]; /* else prev_item = carry (see (A)) OR last tail (see (B)); */ if (update_i < interval_end) { %if is_segmented: if (l_segment_start_flags[LID_0]) prev_item = ${neutral}; %endif scan_type item = ldata[LID_0]; index_type i = update_i; { ${output_statement}; } } if (LID_0 == 0) prev_item = ldata[WG_SIZE - 1]; // (B) local_barrier(); } // }}} %endif } """ # }}} # {{{ driver # {{{ helpers def _round_down_to_power_of_2(val): result = 2**bitlog2(val) if result > val: result >>= 1 assert result <= val return result _PREFIX_WORDS = set(""" ldata partial_scan_buffer global scan_offset segment_start_in_k_group carry g_first_segment_start_in_interval IS_SEG_START tmp Z val l_first_segment_start_in_subtree unit_size index_type interval_begin interval_size offset_end K SCAN_EXPR do_update WG_SIZE first_segment_start_in_k_group scan_type segment_start_in_subtree offset interval_results interval_end first_segment_start_in_subtree unit_base first_segment_start_in_interval k INPUT_EXPR prev_group_sum prev pv value partial_val pgs is_seg_start update_i scan_item_at_i seq_i read_i l_ o_mod_k o_div_k l_segment_start_flags scan_value sum first_seg_start_in_interval g_segment_start_flags group_base seg_end my_val DEBUG ARGS ints_to_store ints_per_wg scan_types_per_int linear_index linear_scan_data_idx dest src store_base wrapped_scan_type dummy scan_tmp tmp_aux LID_2 LID_1 LID_0 LDIM_0 LDIM_1 LDIM_2 GDIM_0 GDIM_1 GDIM_2 GID_0 GID_1 GID_2 """.split()) _IGNORED_WORDS = set(""" 4 8 32 typedef for endfor if void while endwhile endfor endif else const printf None return bool n char true false ifdef pycu_printf str range assert np iinfo max itemsize __packed__ struct __restrict__ extern C set iteritems len setdefault GLOBAL_MEM LOCAL_MEM_ARG WITHIN_KERNEL LOCAL_MEM KERNEL REQD_WG_SIZE local_barrier __syncthreads pragma __attribute__ __global__ __device__ __shared__ __launch_bounds__ threadIdx blockIdx blockDim gridDim x y z barrier _final_update _debug_scan kernel_name positions all padded integer its previous write based writes 0 has local worth scan_expr to read cannot not X items False bank four beginning follows applicable item min each indices works side scanning right summed relative used id out index avoid current state boundary True across be This reads groups along Otherwise undetermined store of times prior s update first regardless Each number because array unit from segment conflicts two parallel 2 empty define direction CL padding work tree bounds values and adds scan is allowed thus it an as enable at in occur sequentially end no storage data 1 largest may representable uses entry Y meaningful computations interval At the left dimension know d A load B group perform shift tail see last OR this add fetched into are directly need gets them stenciled that undefined there up any ones or name only relevant populated even wide we Prepare int seg Note re below place take variable must intra Therefore find code assumption branch workgroup complicated granularity phase remainder than simpler We smaller look ifs lots self behind allow barriers whole loop after already Observe achieve contiguous stores hard go with by math size won t way divisible bit so Avoid declare adding single type is_tail is_first_level input_expr argument_signature preamble double_support neutral output_statement k_group_size name_prefix is_segmented index_dtype scan_dtype wg_size is_segment_start_expr fetch_expr_offsets arg_ctypes ife_offsets input_fetch_exprs def ife_offset arg_name local_fetch_expr_args update_body update_loop_lookbehind update_loop_plain update_loop use_lookbehind_update store_segment_start_flags update_loop first_seg scan_dtype dtype_to_ctype use_bank_conflict_avoidance a b prev_item i last_item prev_value N NO_SEG_BOUNDARY across_seg_boundary """.split()) def _make_template(s): leftovers = set() def replace_id(match): # avoid name clashes with user code by adding 'psc_' prefix to # identifiers. word = match.group(1) if word in _IGNORED_WORDS: return word elif word in _PREFIX_WORDS: return "psc_" + word else: leftovers.add(word) return word import re s = re.sub(r"\b([a-zA-Z0-9_]+)\b", replace_id, s) if leftovers: from warnings import warn warn("leftover words in identifier prefixing: " + " ".join(leftovers)) return mako.template.Template(s, strict_undefined=True) class _GeneratedScanKernelInfo(Record): __slots__ = [ "scan_src", "kernel_name", "scalar_arg_dtypes", "wg_size", "k_group_size"] def __init__(self, scan_src, kernel_name, scalar_arg_dtypes, wg_size, k_group_size): Record.__init__(self, scan_src=scan_src, kernel_name=kernel_name, scalar_arg_dtypes=scalar_arg_dtypes, wg_size=wg_size, k_group_size=k_group_size) def build(self, options): program = SourceModule(self.scan_src, options=options) kernel = program.get_function(self.kernel_name) kernel.prepare(self.scalar_arg_dtypes) return _BuiltScanKernelInfo( kernel=kernel, wg_size=self.wg_size, k_group_size=self.k_group_size) class _BuiltScanKernelInfo(RecordWithoutPickling): __slots__ = ["kernel", "wg_size", "k_group_size"] def __init__(self, kernel, wg_size, k_group_size): RecordWithoutPickling.__init__(self, kernel=kernel, wg_size=wg_size, k_group_size=k_group_size) class _GeneratedFinalUpdateKernelInfo(Record): def __init__(self, source, kernel_name, scalar_arg_dtypes, update_wg_size): Record.__init__(self, source=source, kernel_name=kernel_name, scalar_arg_dtypes=scalar_arg_dtypes, update_wg_size=update_wg_size) def build(self, options): program = SourceModule(self.source, options=options) kernel = program.get_function(self.kernel_name) kernel.prepare(self.scalar_arg_dtypes) return _BuiltFinalUpdateKernelInfo( kernel=kernel, update_wg_size=self.update_wg_size ) class _BuiltFinalUpdateKernelInfo(RecordWithoutPickling): __slots__ = ["kernel", "update_wg_size"] def __init__(self, kernel, update_wg_size): RecordWithoutPickling.__init__(self, kernel=kernel, update_wg_size=update_wg_size) # }}} class ScanPerformanceWarning(UserWarning): pass class _GenericScanKernelBase(object): # {{{ constructor, argument processing def __init__(self, dtype, arguments, input_expr, scan_expr, neutral, output_statement, is_segment_start_expr=None, input_fetch_exprs=[], index_dtype=np.int32, name_prefix="scan", options=None, preamble=""): """ :arg dtype: the :class:`numpy.dtype` with which the scan will be performed. May be a structured type if that type was registered through :func:`pycuda.tools.get_or_register_dtype`. :arg arguments: A string of comma-separated C argument declarations. If *arguments* is specified, then *input_expr* must also be specified. All types used here must be known to PyCUDA. (see :func:`pycuda.tools.get_or_register_dtype`). :arg scan_expr: The associative, binary operation carrying out the scan, represented as a C string. Its two arguments are available as `a` and `b` when it is evaluated. `b` is guaranteed to be the 'element being updated', and `a` is the increment. Thus, if some data is supposed to just propagate along without being modified by the scan, it should live in `b`. This expression may call functions given in the *preamble*. Another value available to this expression is `across_seg_boundary`, a C `bool` indicating whether this scan update is crossing a segment boundary, as defined by `is_segment_start_expr`. The scan routine does not implement segmentation semantics on its own. It relies on `scan_expr` to do this. This value is available (but always `false`) even for a non-segmented scan. .. note:: In early pre-releases of the segmented scan, segmentation semantics were implemented *without* relying on `scan_expr`. :arg input_expr: A C expression, encoded as a string, resulting in the values to which the scan is applied. This may be used to apply a mapping to values stored in *arguments* before being scanned. The result of this expression must match *dtype*. The index intended to be mapped is available as `i` in this expression. This expression may also use the variables defined by *input_fetch_expr*. This expression may also call functions given in the *preamble*. :arg output_statement: a C statement that writes the output of the scan. It has access to the scan result as `item`, the preceding scan result item as `prev_item`, and the current index as `i`. `prev_item` in a segmented scan will be the neutral element at a segment boundary, not the immediately preceding item. Using *prev_item* in output statement has a small run-time cost. `prev_item` enables the construction of an exclusive scan. For non-segmented scans, *output_statement* may also reference `last_item`, which evaluates to the scan result of the last array entry. :arg is_segment_start_expr: A C expression, encoded as a string, resulting in a C `bool` value that determines whether a new scan segments starts at index *i*. If given, makes the scan a segmented scan. Has access to the current index `i`, the result of *input_expr* as a, and in addition may use *arguments* and *input_fetch_expr* variables just like *input_expr*. If it returns true, then previous sums will not spill over into the item with index *i* or subsequent items. :arg input_fetch_exprs: a list of tuples *(NAME, ARG_NAME, OFFSET)*. An entry here has the effect of doing the equivalent of the following before input_expr:: ARG_NAME_TYPE NAME = ARG_NAME[i+OFFSET]; `OFFSET` is allowed to be 0 or -1, and `ARG_NAME_TYPE` is the type of `ARG_NAME`. :arg preamble: |preamble| The first array in the argument list determines the size of the index space over which the scan is carried out, and thus the values over which the index *i* occurring in a number of code fragments in arguments above will vary. All code fragments further have access to N, the number of elements being processed in the scan. """ dtype = self.dtype = np.dtype(dtype) if neutral is None: from warnings import warn warn("not specifying 'neutral' is deprecated and will lead to " "wrong results if your scan is not in-place or your " "'output_statement' does something otherwise non-trivial", stacklevel=2) if dtype.itemsize % 4 != 0: raise TypeError( "scan value type must have size divisible by 4 bytes") self.index_dtype = np.dtype(index_dtype) if np.iinfo(self.index_dtype).min >= 0: raise TypeError("index_dtype must be signed") self.options = options self.parsed_args = parse_arg_list(arguments) from pycuda.tools import VectorArg vector_args_indices = [i for i, arg in enumerate(self.parsed_args) if isinstance(arg, VectorArg)] self.first_array_idx = vector_args_indices[0] self.input_expr = input_expr self.is_segment_start_expr = is_segment_start_expr self.is_segmented = is_segment_start_expr is not None if self.is_segmented: is_segment_start_expr = _process_code_for_macro( is_segment_start_expr) self.output_statement = output_statement for name, arg_name, ife_offset in input_fetch_exprs: if ife_offset not in [0, -1]: raise RuntimeError( "input_fetch_expr offsets must either be 0 or -1") self.input_fetch_exprs = input_fetch_exprs arg_dtypes = {} arg_ctypes = {} for arg in self.parsed_args: arg_dtypes[arg.name] = arg.dtype arg_ctypes[arg.name] = dtype_to_ctype(arg.dtype) self.options = options self.name_prefix = name_prefix # {{{ set up shared code dict from pytools import all from pycuda.characterize import has_double_support self.code_variables = dict( np=np, dtype_to_ctype=dtype_to_ctype, preamble=preamble, name_prefix=name_prefix, index_dtype=self.index_dtype, scan_dtype=dtype, is_segmented=self.is_segmented, arg_dtypes=arg_dtypes, arg_ctypes=arg_ctypes, scan_expr=_process_code_for_macro(scan_expr), neutral=_process_code_for_macro(neutral), double_support=has_double_support(), ) index_typename = dtype_to_ctype(self.index_dtype) scan_typename = dtype_to_ctype(dtype) # This key is meant to uniquely identify the non-device parameters for # the scan kernel. self.kernel_key = ( self.dtype, tuple(arg.declarator() for arg in self.parsed_args), self.input_expr, scan_expr, neutral, output_statement, is_segment_start_expr, tuple(input_fetch_exprs), index_dtype, name_prefix, preamble, # These depend on dtype_to_ctype(), so their value is independent of # the other variables. index_typename, scan_typename, ) # }}} self.use_lookbehind_update = "prev_item" in self.output_statement self.store_segment_start_flags = ( self.is_segmented and self.use_lookbehind_update) self.finish_setup() # }}} generic_scan_kernel_cache = WriteOncePersistentDict( "pycuda-generated-scan-kernel-cache-v1", key_builder=_NumpyTypesKeyBuilder()) class GenericScanKernel(_GenericScanKernelBase): """Generates and executes code that performs prefix sums ("scans") on arbitrary types, with many possible tweaks. Usage example:: import pycuda.gpuarray as gpuarray from compyle.cuda import GenericScanKernel knl = GenericScanKernel( np.int32, arguments="int *ary", input_expr="ary[i]", scan_expr="a+b", neutral="0", output_statement="ary[i] = item;") a = gpuarray.arange(10000, dtype=np.int32) knl(a) """ def finish_setup(self): # Before generating the kernel, see if it's cached. cache_key = (self.kernel_key,) from_cache = False try: result = generic_scan_kernel_cache[cache_key] from_cache = True logger.debug( "cache hit for generated scan kernel '%s'" % self.name_prefix) ( self.first_level_scan_gen_info, self.second_level_scan_gen_info, self.final_update_gen_info) = result except KeyError: pass if not from_cache: logger.debug( "cache miss for generated scan kernel '%s'" % self.name_prefix) self._finish_setup_impl() result = (self.first_level_scan_gen_info, self.second_level_scan_gen_info, self.final_update_gen_info) generic_scan_kernel_cache.store_if_not_present(cache_key, result) # Build the kernels. self.first_level_scan_info = self.first_level_scan_gen_info.build( self.options) del self.first_level_scan_gen_info self.second_level_scan_info = self.second_level_scan_gen_info.build( self.options) del self.second_level_scan_gen_info self.final_update_info = self.final_update_gen_info.build( self.options) del self.final_update_gen_info def _finish_setup_impl(self): # {{{ find usable workgroup/k-group size, build first-level scan trip_count = 0 dev = drv.Context.get_device() avail_local_mem = dev.get_attribute( drv.device_attribute.MAX_SHARED_MEMORY_PER_BLOCK) # not sure where these go, but roughly this much seems unavailable. avail_local_mem -= 0x400 max_scan_wg_size = dev.get_attribute( drv.device_attribute.MAX_THREADS_PER_BLOCK) wg_size_multiples = 64 use_bank_conflict_avoidance = ( self.dtype.itemsize > 4 and self.dtype.itemsize % 8 == 0) # k_group_size should be a power of two because of in-kernel # division by that number. solutions = [] for k_exp in range(0, 9): for wg_size in range(wg_size_multiples, max_scan_wg_size + 1, wg_size_multiples): k_group_size = 2**k_exp lmem_use = self.get_local_mem_use(wg_size, k_group_size, use_bank_conflict_avoidance) if lmem_use <= avail_local_mem: solutions.append( (wg_size * k_group_size, k_group_size, wg_size)) from pytools import any for wg_size_floor in [256, 192, 128]: have_sol_above_floor = any(wg_size >= wg_size_floor for _, _, wg_size in solutions) if have_sol_above_floor: # delete all solutions not meeting the wg size floor solutions = [(total, try_k_group_size, try_wg_size) for total, try_k_group_size, try_wg_size in solutions if try_wg_size >= wg_size_floor] break _, k_group_size, max_scan_wg_size = max(solutions) while True: candidate_scan_gen_info = self.generate_scan_kernel( max_scan_wg_size, self.parsed_args, _process_code_for_macro(self.input_expr), self.is_segment_start_expr, input_fetch_exprs=self.input_fetch_exprs, is_first_level=True, store_segment_start_flags=self.store_segment_start_flags, k_group_size=k_group_size, use_bank_conflict_avoidance=use_bank_conflict_avoidance) candidate_scan_info = candidate_scan_gen_info.build( self.options) # Will this device actually let us execute this kernel # at the desired work group size? Building it is the # only way to find out. kernel_max_wg_size = candidate_scan_info.kernel.get_attribute( drv.function_attribute.MAX_THREADS_PER_BLOCK) if candidate_scan_info.wg_size <= kernel_max_wg_size: break else: max_scan_wg_size = min(kernel_max_wg_size, max_scan_wg_size) trip_count += 1 assert trip_count <= 20 self.first_level_scan_gen_info = candidate_scan_gen_info assert (_round_down_to_power_of_2(candidate_scan_info.wg_size) == candidate_scan_info.wg_size) # }}} # {{{ build second-level scan from pycuda.tools import VectorArg second_level_arguments = self.parsed_args + [ VectorArg(self.dtype, "interval_sums")] second_level_build_kwargs = {} if self.is_segmented: second_level_arguments.append( VectorArg(self.index_dtype, "g_first_segment_start_in_interval_input")) # is_segment_start_expr answers the question "should previous sums # spill over into this item". And since # g_first_segment_start_in_interval_input answers the question if a # segment boundary was found in an interval of data, then if not, # it's ok to spill over. second_level_build_kwargs["is_segment_start_expr"] = \ "g_first_segment_start_in_interval_input[i] != NO_SEG_BOUNDARY" else: second_level_build_kwargs["is_segment_start_expr"] = None self.second_level_scan_gen_info = self.generate_scan_kernel( max_scan_wg_size, arguments=second_level_arguments, input_expr="interval_sums[i]", input_fetch_exprs=[], is_first_level=False, store_segment_start_flags=False, k_group_size=k_group_size, use_bank_conflict_avoidance=use_bank_conflict_avoidance, **second_level_build_kwargs) # }}} # {{{ generate final update kernel update_wg_size = min(max_scan_wg_size, 256) final_update_tpl = _make_template(UPDATE_SOURCE) final_update_src = str(final_update_tpl.render( wg_size=update_wg_size, output_statement=self.output_statement, argument_signature=", ".join( arg.declarator() for arg in self.parsed_args), is_segment_start_expr=self.is_segment_start_expr, input_expr=_process_code_for_macro(self.input_expr), use_lookbehind_update=self.use_lookbehind_update, **self.code_variables)) update_scalar_arg_dtypes = ( get_arg_list_scalar_arg_dtypes(self.parsed_args) + [self.index_dtype, self.index_dtype, None, None]) if self.is_segmented: # g_first_segment_start_in_interval update_scalar_arg_dtypes.append(None) if self.store_segment_start_flags: update_scalar_arg_dtypes.append(None) # g_segment_start_flags self.final_update_gen_info = _GeneratedFinalUpdateKernelInfo( final_update_src, self.name_prefix + "_final_update", update_scalar_arg_dtypes, update_wg_size) # }}} # {{{ scan kernel build/properties def get_local_mem_use(self, k_group_size, wg_size, use_bank_conflict_avoidance): arg_dtypes = {} for arg in self.parsed_args: arg_dtypes[arg.name] = arg.dtype fetch_expr_offsets = {} for name, arg_name, ife_offset in self.input_fetch_exprs: fetch_expr_offsets.setdefault(arg_name, set()).add(ife_offset) itemsize = self.dtype.itemsize if use_bank_conflict_avoidance: itemsize += 4 return ( # ldata itemsize * (k_group_size + 1) * (wg_size + 1) # l_segment_start_flags + k_group_size * wg_size # l_first_segment_start_in_subtree + self.index_dtype.itemsize * wg_size + k_group_size * wg_size * sum( arg_dtypes[arg_name].itemsize for arg_name, ife_offsets in list(fetch_expr_offsets.items()) if -1 in ife_offsets or len(ife_offsets) > 1)) def generate_scan_kernel( self, max_wg_size, arguments, input_expr, is_segment_start_expr, input_fetch_exprs, is_first_level, store_segment_start_flags, k_group_size, use_bank_conflict_avoidance): scalar_arg_dtypes = get_arg_list_scalar_arg_dtypes(arguments) # Empirically found on Nv hardware: no need to be bigger than this size wg_size = _round_down_to_power_of_2( min(max_wg_size, 256)) kernel_name = self.code_variables["name_prefix"] if is_first_level: kernel_name += "_lev1" else: kernel_name += "_lev2" scan_tpl = _make_template(SCAN_INTERVALS_SOURCE) scan_src = str( scan_tpl.render( wg_size=wg_size, input_expr=input_expr, k_group_size=k_group_size, argument_signature=", ".join( arg.declarator() for arg in arguments), is_segment_start_expr=is_segment_start_expr, input_fetch_exprs=input_fetch_exprs, is_first_level=is_first_level, store_segment_start_flags=store_segment_start_flags, use_bank_conflict_avoidance=use_bank_conflict_avoidance, kernel_name=kernel_name, **self.code_variables)) scalar_arg_dtypes.extend( (None, self.index_dtype, self.index_dtype)) if is_first_level: scalar_arg_dtypes.append(None) # interval_results if self.is_segmented and is_first_level: scalar_arg_dtypes.append(None) # g_first_segment_start_in_interval if store_segment_start_flags: scalar_arg_dtypes.append(None) # g_segment_start_flags return _GeneratedScanKernelInfo( scan_src=scan_src, kernel_name=kernel_name, scalar_arg_dtypes=scalar_arg_dtypes, wg_size=wg_size, k_group_size=k_group_size) # }}} def __call__(self, *args, **kwargs): # {{{ argument processing allocator = kwargs.get("allocator") n = kwargs.get("size") stream = kwargs.get("stream") if len(args) != len(self.parsed_args): raise TypeError("expected %d arguments, got %d" % (len(self.parsed_args), len(args))) first_array = args[self.first_array_idx] allocator = allocator or first_array.allocator if n is None: n, = first_array.shape if n == 0: return data_args = [] from pycuda.tools import VectorArg for arg_descr, arg_val in zip(self.parsed_args, args): if isinstance(arg_descr, VectorArg): data_args.append(arg_val.gpudata) else: data_args.append(arg_val) # }}} l1_info = self.first_level_scan_info l2_info = self.second_level_scan_info unit_size = l1_info.wg_size * l1_info.k_group_size dev = drv.Context.get_device() max_intervals = 3 * dev.get_attribute( drv.device_attribute.MULTIPROCESSOR_COUNT) from pytools import uniform_interval_splitting interval_size, num_intervals = uniform_interval_splitting( n, unit_size, max_intervals) # {{{ allocate some buffers interval_results = gpuarray.empty( num_intervals, dtype=self.dtype, allocator=allocator) partial_scan_buffer = gpuarray.empty( n, dtype=self.dtype, allocator=allocator) if self.store_segment_start_flags: segment_start_flags = gpuarray.empty( n, dtype=bool, allocator=allocator) # }}} # {{{ first level scan of interval (one interval per block) scan1_args = data_args + [ partial_scan_buffer.gpudata, n, interval_size, interval_results.gpudata, ] if self.is_segmented: first_segment_start_in_interval = gpuarray.empty( num_intervals, dtype=self.index_dtype, allocator=allocator) scan1_args.append(first_segment_start_in_interval.gpudata) if self.store_segment_start_flags: scan1_args.append(segment_start_flags.gpudata) l1_evt = l1_info.kernel.prepared_async_call( (num_intervals, 1), (l1_info.wg_size, 1, 1), stream, *scan1_args) # }}} # {{{ second level scan of per-interval results # can scan at most one interval assert interval_size >= num_intervals scan2_args = data_args + [ interval_results.gpudata, # interval_sums ] if self.is_segmented: scan2_args.append(first_segment_start_in_interval.gpudata) scan2_args = scan2_args + [ interval_results.gpudata, # partial_scan_buffer num_intervals, interval_size] l2_evt = l2_info.kernel.prepared_async_call( (1, 1), (l1_info.wg_size, 1, 1), stream, *scan2_args) # }}} # {{{ update intervals with result of interval scan upd_args = data_args + [n, interval_size, interval_results.gpudata, partial_scan_buffer.gpudata] if self.is_segmented: upd_args.append(first_segment_start_in_interval.gpudata) if self.store_segment_start_flags: upd_args.append(segment_start_flags.gpudata) return self.final_update_info.kernel.prepared_async_call( (num_intervals, 1), (self.final_update_info.update_wg_size, 1, 1), stream, *upd_args) # }}} # }}}