<s> import numpy as np <EOL> import gpustats . kernels as kernels <EOL> import gpustats . codegen as codegen <EOL> import gpustats . util as util <EOL> import pycuda . driver as drv <EOL> from pycuda . gpuarray import GPUArray , to_gpu <EOL> from pycuda . gpuarray import empty as gpu_empty <EOL> from pycuda . curandom import rand as curand <EOL> cu_module = codegen . get_full_cuda_module ( ) <EOL> def sample_discrete ( densities , logged = False , <EOL> return_gpuarray = False ) : <EOL> """<STR_LIT>""" <EOL> from gpustats . util import info <EOL> n , k = densities . shape <EOL> if isinstance ( densities , GPUArray ) : <EOL> if densities . flags . f_contiguous : <EOL> gpu_densities = util . transpose ( densities ) <EOL> else : <EOL> gpu_densities = densities <EOL> else : <EOL> densities = util . prep_ndarray ( densities ) <EOL> gpu_densities = to_gpu ( densities ) <EOL> cu_func = cu_module . get_function ( '<STR_LIT>' ) <EOL> gpu_random = to_gpu ( np . asarray ( np . random . rand ( n ) , dtype = np . float32 ) ) <EOL> gpu_dest = gpu_empty ( n , dtype = np . int32 ) <EOL> dims = np . array ( [ n , k , logged ] , dtype = np . int32 ) <EOL> if info . max_block_threads < <NUM_LIT> : <EOL> x_block_dim = <NUM_LIT:16> <EOL> else : <EOL> x_block_dim = <NUM_LIT:32> <EOL> y_block_dim = <NUM_LIT:16> <EOL> block_design = ( x_block_dim , y_block_dim , <NUM_LIT:1> ) <EOL> grid_design = ( int ( n / y_block_dim ) + <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> shared_mem = <NUM_LIT:4> * ( ( x_block_dim + <NUM_LIT:1> ) * y_block_dim + <EOL> <NUM_LIT:2> * y_block_dim ) <EOL> cu_func ( gpu_densities , gpu_random , gpu_dest , <EOL> dims [ <NUM_LIT:0> ] , dims [ <NUM_LIT:1> ] , dims [ <NUM_LIT:2> ] , <EOL> block = block_design , grid = grid_design , shared = shared_mem ) <EOL> gpu_random . gpudata . free ( ) <EOL> if return_gpuarray : <EOL> return gpu_dest <EOL> else : <EOL> res = gpu_dest . get ( ) <EOL> gpu_dest . gpudata . free ( ) <EOL> return res <EOL> def sample_discrete_old ( in_densities , logged = False , pad = False , <EOL> return_gpuarray = False ) : <EOL> """<STR_LIT>""" <EOL> if pad : <EOL> if logged : <EOL> densities = util . pad_data_mult16 ( in_densities , fill = <NUM_LIT:1> ) <EOL> else : <EOL> densities = util . pad_data_mult16 ( in_densities , fill = <NUM_LIT:0> ) <EOL> else : <EOL> densities = in_densities <EOL> n , k = densities . shape <EOL> if logged : <EOL> cu_func = cu_module . get_function ( '<STR_LIT>' ) <EOL> else : <EOL> cu_func = cu_module . get_function ( '<STR_LIT>' ) <EOL> if isinstance ( densities , GPUArray ) : <EOL> if densities . flags . f_contiguous : <EOL> gpu_densities = util . transpose ( densities ) <EOL> else : <EOL> gpu_densities = densities <EOL> else : <EOL> densities = util . prep_ndarray ( densities ) <EOL> gpu_densities = to_gpu ( densities ) <EOL> gpu_random = to_gpu ( np . asarray ( np . random . rand ( n ) , dtype = np . float32 ) ) <EOL> gpu_dest = gpu_empty ( n , dtype = np . float32 ) <EOL> stride = gpu_densities . shape [ <NUM_LIT:1> ] <EOL> if stride % <NUM_LIT:2> == <NUM_LIT:0> : <EOL> stride += <NUM_LIT:1> <EOL> dims = np . array ( [ n , k , gpu_densities . shape [ <NUM_LIT:1> ] , stride ] , dtype = np . int32 ) <EOL> grid_design , block_design = _tune_sfm ( n , stride , cu_func . num_regs ) <EOL> shared_mem = <NUM_LIT:4> * ( block_design [ <NUM_LIT:0> ] * stride + <EOL> <NUM_LIT:1> * block_design [ <NUM_LIT:0> ] ) <EOL> cu_func ( gpu_densities , gpu_random , gpu_dest , <EOL> dims [ <NUM_LIT:0> ] , dims [ <NUM_LIT:1> ] , dims [ <NUM_LIT:2> ] , dims [ <NUM_LIT:3> ] , <EOL> block = block_design , grid = grid_design , shared = shared_mem ) <EOL> gpu_random . gpudata . free ( ) <EOL> if return_gpuarray : <EOL> return gpu_dest <EOL> else : <EOL> res = gpu_dest . get ( ) <EOL> gpu_dest . gpudata . free ( ) <EOL> return res <EOL> def _tune_sfm ( n , stride , func_regs ) : <EOL> """<STR_LIT>""" <EOL> from gpustats . util import info <EOL> comp_cap = info . compute_cap <EOL> max_smem = info . shared_mem * <NUM_LIT> <EOL> max_threads = int ( info . max_block_threads * <NUM_LIT:0.5> ) <EOL> max_regs = <NUM_LIT> * info . max_registers <EOL> if comp_cap [ <NUM_LIT:0> ] == <NUM_LIT:1> : <EOL> xdim = <NUM_LIT:16> <EOL> else : <EOL> xdim = <NUM_LIT:32> <EOL> def sfm_config_ok ( xdim , ydim , stride , func_regs , max_regs , max_smem , max_threads ) : <EOL> ok = <NUM_LIT:4> * ( xdim * stride + <NUM_LIT:1> * xdim ) < max_smem and func_regs * ydim * xdim < max_regs <EOL> return ok and xdim * ydim <= max_threads <EOL> ydim = <NUM_LIT:2> <EOL> while sfm_config_ok ( xdim , ydim , stride , func_regs , max_regs , max_smem , max_threads ) : <EOL> ydim += <NUM_LIT:1> <EOL> ydim -= <NUM_LIT:1> <EOL> nblocks = int ( n / xdim ) + <NUM_LIT:1> <EOL> return ( nblocks , <NUM_LIT:1> ) , ( xdim , ydim , <NUM_LIT:1> ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> n = <NUM_LIT:100> <EOL> k = <NUM_LIT:5> <EOL> dens = np . log ( np . abs ( np . random . randn ( k ) ) ) - <NUM_LIT:200> <EOL> densities = [ dens . copy ( ) for _ in range ( n ) ] <EOL> dens = np . exp ( dens + <NUM_LIT:200> ) <EOL> densities = np . asarray ( densities ) <EOL> labels = sample_discrete ( densities , logged = True ) <EOL> mu = np . dot ( dens / dens . sum ( ) , np . arange ( k ) ) <EOL> print mu , labels . mean ( ) </s>
