def _configure_template_vals_bprop(in_dtype, dtype): template_vals = { "in_dtype": _get_register_type(in_dtype, memory=True), "type": _get_register_type(dtype, memory=True), "compute_dW_code": r"""dW[output_row + i] += errors[error_row + i];""", "stats_args": "", "atomic_max": "", "compute_input": "", "common": "", } if isinstance(dtype, Flex): template_vals[ "stats_args"] = ", int* maxabs, float scaleO, float scaleI, float scaleE" template_vals[ "atomic_max"] = r"""atomicMax(maxabs, intermediate_max);""" template_vals["compute_input"] = "* scaleI" template_vals["common"] = _flex_includes_template template_vals["compute_dW_code"] = r""" float dW_float = dW[output_row + i] * scaleO; float error_float = errors[error_row + i] * scaleE; dW_float += error_float; int dW_out = fp32_to_int16(dW_float / scaleO); intermediate_max = max_abs(intermediate_max, dW_out); dW[output_row + i] = dW_out; """ return template_vals
def _get_transpose_kernel(dtype): code = _transpose_kernel % {"type": _get_register_type(dtype, memory=True)} module = SourceModule(code) kernel = module.get_function("transpose") kernel.prepare("PPII") return kernel
def _get_oned_copy_kernel(dtype, shape): copy = r""" __global__ void copy_oned(%(type)s* out, const %(type)s* in, int dim, long long src_str, long long dst_str) { int tid_x = threadIdx.x; int idx = blockIdx.x; idx = (idx << 5) + tid_x; const %(type)s* in0 = in + (src_str * idx); %(type)s* out0 = out + (dst_str * idx); if(idx < dim) *out0 = *in0; } """ code = copy % dict( type=_get_register_type(dtype, memory=True) ) # print code module = SourceModule(code) kernel = module.get_function("copy_oned") kernel.prepare("PPIqq") kernel.grid = (_ceil_div(shape[0], 32), 1, 1) kernel.block = (32, 1, 1) kernel.args = (shape[0], ) return kernel
def _configure_template_vals_sort(block_size, kernel_id, in_dtype): template_vals = { "threads": block_size, "store_blocksum": (1 if kernel_id == 1 else 0), "in_dtype": _get_register_type(in_dtype, memory=True), "stats_args": "", "compute_input": "", } if isinstance(in_dtype, Flex): template_vals["stats_args"] = ", float scaleI" template_vals["compute_input"] = "* scaleI" return template_vals
def _get_lut_bprop_kernel(dtype, in_dtype, deterministic=False): """ Builds the bprop kernel for lookup table layers based on templated code. If the deterministic version is requested, an index buffer must be passed as an argument. This index buffer re-orders items in the input tensor so that word_ids are sorted. This is required since we need to be sure that each thread only updates weights for one word id. Arguments: dtype (np.dtype): The data which the kernel will operate on. deterministic (boolean): Builds the deterministic kernel when this is set to True. """ if not deterministic: code = r""" __global__ void lut_bprop( %(in_dtype)s* inputs, %(type)s* dW, %(type)s* errors, const int nin, const int embedding_dim, const int vocab_size, const int pad_idx) { const int tid = threadIdx.x; const int bid = blockIdx.x; int word_id = inputs[bid]; int error_row = bid * embedding_dim; int output_row = word_id * embedding_dim; if(word_id != pad_idx) { for(int i = tid; i < embedding_dim; i += blockDim.x) { atomicAdd(&dW[output_row + i], errors[error_row + i]); } } } """ code = code % {"type": _get_register_type(dtype)} module = SourceModule(code, options=["--use_fast_math"]) kernel = module.get_function("lut_bprop") kernel.prepare("PPPIIIi") else: code = r""" %(common)s __global__ void lut_bprop( %(in_dtype)s* inputs, int* index_buffer, %(type)s* dW, %(type)s* errors, const int nin, const int embedding_dim, const int vocab_size, const int pad_idx %(stats_args)s) { const int tid = threadIdx.x; const int bid = blockIdx.x; int index_position = bid; int index = index_buffer[index_position]; int word_id = inputs[index] %(compute_input)s; int intermediate_max = 0; if((bid == 0 || word_id != (inputs[index_buffer[bid - 1]] %(compute_input)s)) && word_id != pad_idx) { int output_row = word_id * embedding_dim; do { int error_row = index * embedding_dim; for(int i = tid; i < embedding_dim; i += blockDim.x) { %(compute_dW_code)s } index_position++; if(index_position == gridDim.x) { break; } index = index_buffer[index_position]; } while((inputs[index] %(compute_input)s) == word_id); } %(atomic_max)s } """ code %= _configure_template_vals_bprop(in_dtype, dtype) module = SourceModule(code, options=["--use_fast_math"]) kernel = module.get_function("lut_bprop") kernel.prepare("PPPPIIIi" + flex_sig_bprop(in_dtype.str[1])) kernel.name = lut_bprop_kernel_name return kernel
def _get_sorting_kernel(kernel_id, block_size, in_dtype): """ Builds kernels used for sorting inputs. There are several kernels here corresponding to the steps in the algorithm. The algorithm works by determining the sorted position for each input item. This is done with a bucket sort algorithm, where each word_id is a bucket. The first step determines the size of each bucket (number of occurences of each word_id). Next, a prefix some is computed over the list of bucket sizes to find where each bucket will be placed in the output buffer. Finally, each thread places it's index into the correct sorted position based on the bucket start index (computed from the prefix sum) and that thread's offset into the bucket (which is taken from the output of the atomic add done in the first step.) Arguments: kernel_id (Integer): Which step to build the kernel for [0, 4] block_size (Integer): Number of threads per block for the prefix sum kernels. """ code = r""" #define THREADS %(threads)s #define STORE_BLOCKSUM %(store_blocksum)s __global__ void sort_inputs0( %(in_dtype)s* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size, const int input_length) { const int tid = threadIdx.x + (blockDim.x * blockIdx.x); int word_id; if(tid < input_length) { word_id = inputs[tid]; offset_buffer[tid] = atomicAdd(&word_counts[word_id], 1); } } __device__ void scan(int* buffer, int* blocksum, int global_length) { const int tid = (threadIdx.x << 1) + 1; const int gid = ((threadIdx.x + (blockIdx.x * blockDim.x)) << 1) + 1; __shared__ int local_counts[THREADS * 2]; local_counts[tid] = buffer[gid]; local_counts[tid - 1] = buffer[gid - 1]; #pragma unroll for(int skip = 1; skip <= THREADS; skip <<= 1) { int mask = (skip << 1) - 1; if((tid & mask) == mask) { local_counts[tid] += local_counts[tid - skip]; } __syncthreads(); } if(tid == (THREADS * 2 - 1)) { #if STORE_BLOCKSUM blocksum[blockIdx.x] = local_counts[tid]; #endif local_counts[tid] = 0; } #pragma unroll for(int skip = THREADS; skip > 0; skip >>= 1) { int mask = (skip << 1) - 1; if((tid & mask) == mask) { int temp = local_counts[tid - skip]; local_counts[tid - skip] = local_counts[tid]; local_counts[tid] += temp; } __syncthreads(); } if(gid < global_length) { buffer[gid] = local_counts[tid]; buffer[gid - 1] = local_counts[tid - 1]; } } __global__ void sort_inputs1( %(in_dtype)s* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size, const int input_length) { scan(word_counts, word_counts + vocab_size, vocab_size); } __global__ void sort_inputs2( %(in_dtype)s* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size, const int input_length) { scan(word_counts + vocab_size, 0, blockDim.x); } __global__ void sort_inputs3( %(in_dtype)s* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size, const int input_length) { const int gid = (threadIdx.x + (blockIdx.x * blockDim.x)) << 1; if(gid < vocab_size) { word_counts[gid] += word_counts[vocab_size + blockIdx.x]; word_counts[gid + 1] += word_counts[vocab_size + blockIdx.x]; } } __global__ void sort_inputs4( %(in_dtype)s* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size, const int input_length) { const int tid = threadIdx.x + (blockDim.x * blockIdx.x); int word_id; if(tid < input_length) { word_id = inputs[tid]; int sorted_position = word_counts[word_id] + offset_buffer[tid]; index_buffer[sorted_position] = tid; } } """ code = code % { "threads": block_size, "store_blocksum": (1 if kernel_id == 1 else 0), "in_dtype": _get_register_type(in_dtype) } module = SourceModule(code, options=["--use_fast_math"]) function_name = "sort_inputs" + native_str(kernel_id) kernel = module.get_function(function_name) kernel.prepare("PPPPII") kernel.name = "sort_inputs" return kernel
def _get_copy_transpose_kernel(dtype, shape, axes=None): if len(shape) == 1: return _get_oned_copy_kernel(dtype, shape) src = list(range(len(shape))) dst = list(axes) src_dim = src[-1] dst_dim = dst[-1] # If the inner dim is the same for both, no need for shared memory tile # Then map the outer source dim to the threadIdx.y values if src_dim == dst_dim: dst_dim = src[0] shared_tile = False else: shared_tile = True src_offset = [] dst_offset = [] params = [] values = [] magic = "" # add dims for bounds checking for dim in (src_dim, dst_dim): params.append("int dim_%s" % dim) values.append(shape[dim]) # collapse src and dst shape by 32 grid_shape = list(shape) grid_shape[src_dim] = _ceil_div(shape[src_dim], 32) grid_shape[dst_dim] = _ceil_div(shape[dst_dim], 32) # get a src list without dst dim src2 = [s for s in src if s != dst_dim] # get the name of the first compound index blkx_name = compound_idx = "".join(native_str(x) for x in src2) # generate the magic number math to extract all indeces while len(src2) > 1: idx1 = src2[0] del src2[0] idx2 = "".join(native_str(i) for i in src2) div = reduce(mul, (grid_shape[i] for i in src2), 1) params.extend(p % idx2 for p in ("int magic_%s", "int shift_%s", "int div_%s")) values.extend(_magic64(div)) values.append(div) magic += r""" int idx_{1} = div64(idx_{0}, magic_{2}, shift_{2}); int idx_{2} = idx_{0} - idx_{1}*div_{2}; """.format(compound_idx, idx1, idx2) compound_idx = idx2 # Add params for src strides and generate src offset # The param values will be added externally for s in src: params.append("long long src_str_%d" % s) src_offset.append("src_str_%d*idx_%d" % (s, s)) # Add params for dst strides and generate dst offset for d in dst: params.append("long long dst_str_%d" % d) dst_offset.append("dst_str_%d*idx_%d" % (d, d)) num_strides = len(src) + len(dst) if shared_tile: copy_transpose = r""" %(common)s __global__ void copy_transpose(%(type)s* out, const %(type)s* in, %(params)s) { __shared__ %(type)s tile[32][33]; int tid_x = threadIdx.x; int tid_y = threadIdx.y; int idx_%(blk)s = blockIdx.x; int idx_%(dst)s = blockIdx.y; %(magic)s idx_%(src)s = (idx_%(src)s << 5) + tid_x; idx_%(dst)s = (idx_%(dst)s << 5) + tid_y; const %(type)s* in00 = in + %(src_offset)s; const %(type)s* in08 = in00 + src_str_%(dst)s*8; const %(type)s* in16 = in08 + src_str_%(dst)s*8; const %(type)s* in24 = in16 + src_str_%(dst)s*8; bool b%(src)s = idx_%(src)s < dim_%(src)s; if (idx_%(dst)s + 0 < dim_%(dst)s && b%(src)s) tile[tid_y + 0][tid_x] = *in00; if (idx_%(dst)s + 8 < dim_%(dst)s && b%(src)s) tile[tid_y + 8][tid_x] = *in08; if (idx_%(dst)s + 16 < dim_%(dst)s && b%(src)s) tile[tid_y + 16][tid_x] = *in16; if (idx_%(dst)s + 24 < dim_%(dst)s && b%(src)s) tile[tid_y + 24][tid_x] = *in24; __syncthreads(); %(type)s val00 = tile[tid_x][tid_y + 0]; %(type)s val08 = tile[tid_x][tid_y + 8]; %(type)s val16 = tile[tid_x][tid_y + 16]; %(type)s val24 = tile[tid_x][tid_y + 24]; idx_%(src)s += tid_y - tid_x; idx_%(dst)s += tid_x - tid_y; bool b%(dst)s = idx_%(dst)s < dim_%(dst)s; %(type)s* out00 = out + %(dst_offset)s; %(type)s* out08 = out00 + dst_str_%(src)s*8; %(type)s* out16 = out08 + dst_str_%(src)s*8; %(type)s* out24 = out16 + dst_str_%(src)s*8; if (idx_%(src)s + 0 < dim_%(src)s && b%(dst)s) *out00 = val00; if (idx_%(src)s + 8 < dim_%(src)s && b%(dst)s) *out08 = val08; if (idx_%(src)s + 16 < dim_%(src)s && b%(dst)s) *out16 = val16; if (idx_%(src)s + 24 < dim_%(src)s && b%(dst)s) *out24 = val24; } """ else: copy_transpose = r""" %(common)s __global__ void copy_transpose(%(type)s* out, const %(type)s* in, %(params)s) { int tid_x = threadIdx.x; int tid_y = threadIdx.y; int idx_%(blk)s = blockIdx.x; int idx_%(dst)s = blockIdx.y; %(magic)s idx_%(src)s = (idx_%(src)s << 5) + tid_x; idx_%(dst)s = (idx_%(dst)s << 5) + tid_y; bool b%(src)s = idx_%(src)s < dim_%(src)s; bool b%(dst)s_00 = idx_%(dst)s + 0 < dim_%(dst)s && b%(src)s; bool b%(dst)s_08 = idx_%(dst)s + 8 < dim_%(dst)s && b%(src)s; bool b%(dst)s_16 = idx_%(dst)s + 16 < dim_%(dst)s && b%(src)s; bool b%(dst)s_24 = idx_%(dst)s + 24 < dim_%(dst)s && b%(src)s; %(type)s val00 = 0; %(type)s val08 = 0; %(type)s val16 = 0; %(type)s val24 = 0; const %(type)s* in00 = in + %(src_offset)s; const %(type)s* in08 = in00 + src_str_%(dst)s*8; const %(type)s* in16 = in08 + src_str_%(dst)s*8; const %(type)s* in24 = in16 + src_str_%(dst)s*8; if (b%(dst)s_00) val00 = *in00; if (b%(dst)s_08) val08 = *in08; if (b%(dst)s_16) val16 = *in16; if (b%(dst)s_24) val24 = *in24; %(type)s* out00 = out + %(dst_offset)s; %(type)s* out08 = out00 + dst_str_%(dst)s*8; %(type)s* out16 = out08 + dst_str_%(dst)s*8; %(type)s* out24 = out16 + dst_str_%(dst)s*8; if (b%(dst)s_00) *out00 = val00; if (b%(dst)s_08) *out08 = val08; if (b%(dst)s_16) *out16 = val16; if (b%(dst)s_24) *out24 = val24; } """ code = copy_transpose % dict( common=_div64, type=_get_register_type(dtype, memory=True), params=", ".join(params), blk=blkx_name, src=src_dim, dst=dst_dim, magic=magic, src_offset=" + ".join(src_offset), dst_offset=" + ".join(dst_offset) ) # print code module = SourceModule(code) kernel = module.get_function("copy_transpose") kernel.prepare("PP" + ("I" * (len(params) - num_strides)) + "q" * num_strides) grid_x = grid_shape[src_dim] grid_y = grid_shape[dst_dim] for s in src: if s not in (src_dim, dst_dim): grid_x *= grid_shape[s] kernel.grid = (grid_x, grid_y, 1) kernel.block = (32, 8, 1) kernel.args = tuple(values) return kernel