codepy.elementwise

1 """Elementwise functionality.""" 2 3 from __future__ import division 4 5 __copyright__ = "Copyright (C) 2009 Andreas Kloeckner" 6 7 8 9 10 from pytools import memoize 11 import numpy 12 from codepy.cgen import POD, Value, dtype_to_ctype

13 14 15 16 17 -class Argument:

18 - def __init__(self, dtype, name):

19 self.dtype = numpy.dtype(dtype) 20 self.name = name

21

22 - def __repr__(self):

23 return "%s(%r, %s)" % ( 24 self.__class__.__name__, 25 self.name, 26 self.dtype)

27

28 -class VectorArg(Argument):

29 - def declarator(self):

30 return Value("numpy_array<%s >" % dtype_to_ctype(self.dtype), 31 self.name+"_ary")

32

33 - def arg_name(self):

34 return self.name + "_ary"

35 36 struct_char = "P"

37

38 -class ScalarArg(Argument):

39 - def declarator(self):

40 return POD(self.dtype, self.name)

41

42 - def arg_name(self):

43 return self.name

44 45 @property

46 - def struct_char(self):

47 return self.dtype.char

48

49 50 51 52 -def get_elwise_module_descriptor(arguments, operation, name="kernel"):

53 from codepy.bpl import BoostPythonModule 54 55 from codepy.cgen import FunctionBody, FunctionDeclaration, \ 56 Value, POD, Struct, For, Initializer, Include, Statement, \ 57 Line, Block 58 59 S = Statement 60 61 mod = BoostPythonModule() 62 mod.add_to_preamble([ 63 Include("pyublas/numpy.hpp"), 64 ]) 65 66 mod.add_to_module([ 67 S("namespace ublas = boost::numeric::ublas"), 68 S("using namespace pyublas"), 69 Line(), 70 ]) 71 72 body = Block([ 73 Initializer( 74 Value("numpy_array<%s >::iterator" 75 % dtype_to_ctype(varg.dtype), 76 varg.name), 77 "args.%s_ary.begin()" % varg.name) 78 for varg in arguments if isinstance(varg, VectorArg)] 79 +[Initializer( 80 sarg.declarator(), "args." + sarg.name) 81 for sarg in arguments if isinstance(sarg, ScalarArg)] 82 ) 83 84 body.extend([ 85 Line(), 86 For("unsigned i = 0", 87 "i < codepy_length", 88 "++i", 89 Block([S(operation)]) 90 ) 91 ]) 92 93 arg_struct = Struct("arg_struct", 94 [arg.declarator() for arg in arguments]) 95 mod.add_struct(arg_struct, "ArgStruct") 96 mod.add_to_module([Line()]) 97 98 mod.add_function( 99 FunctionBody( 100 FunctionDeclaration( 101 Value("void", name), 102 [POD(numpy.uintp, "codepy_length"), 103 Value("arg_struct", "args")]), 104 body)) 105 106 return mod

107

108 109 110 -def get_elwise_module_binary(arguments, operation, name="kernel", toolchain=None, 111 wait_on_error=False):

112 if toolchain is None: 113 from codepy.jit import guess_toolchain 114 toolchain = guess_toolchain() 115 116 toolchain = toolchain.copy() 117 118 from codepy.libraries import add_pyublas 119 toolchain = toolchain.copy() 120 add_pyublas(toolchain) 121 122 return get_elwise_module_descriptor(arguments, operation, name) \ 123 .compile(toolchain, wait_on_error=wait_on_error)

124

125 126 127 128 -def get_elwise_kernel(arguments, operation, name="kernel", toolchain=None, 129 wait_on_error=False):

130 return getattr(get_elwise_module_binary( 131 arguments, operation, name, toolchain, wait_on_error), name)

132

133 134 135 136 -class ElementwiseKernel:

137 - def __init__(self, arguments, operation, name="kernel", toolchain=None, 138 wait_on_error=False):

139 self.arguments = arguments 140 self.module = get_elwise_module_binary( 141 arguments, operation, name, toolchain, wait_on_error) 142 self.func = getattr(self.module, name) 143 144 self.vec_arg_indices = [i for i, arg in enumerate(arguments) 145 if isinstance(arg, VectorArg)] 146 147 assert self.vec_arg_indices, \ 148 "ElementwiseKernel can only be used with functions that have at least one " \ 149 "vector argument"

150

151 - def __call__(self, *args):

152 vectors = [] 153 154 from pytools import single_valued 155 size = single_valued(args[i].size for i in self.vec_arg_indices) 156 157 # no need to do type checking--pyublas does that for us 158 arg_struct = self.module.ArgStruct() 159 for arg_descr, arg in zip(self.arguments, args): 160 setattr(arg_struct, arg_descr.arg_name(), arg) 161 162 assert not arg_struct.__dict__ 163 164 self.func(size, arg_struct)

165

166 167 168 169 @memoize 170 -def make_linear_comb_kernel_with_result_dtype( 171 result_dtype, scalar_dtypes, vector_dtypes):

172 comp_count = len(vector_dtypes) 173 from pytools import flatten 174 return ElementwiseKernel([VectorArg(result_dtype, "result")] + list(flatten( 175 (ScalarArg(scalar_dtypes[i], "a%d_fac" % i), 176 VectorArg(vector_dtypes[i], "a%d" % i)) 177 for i in range(comp_count))), 178 "result[i] = " + " + ".join("a%d_fac*a%d[i]" % (i, i) 179 for i in range(comp_count)))

180

181 182 183 184 @memoize 185 -def make_linear_comb_kernel(scalar_dtypes, vector_dtypes):

186 from pytools import common_dtype 187 result_dtype = common_dtype(scalar_dtypes+vector_dtypes) 188 189 return make_linear_comb_kernel_with_result_dtype( 190 result_dtype, scalar_dtypes, vector_dtypes), result_dtype

191 192 193 194 195 196 if __name__ == "__main__": 197 import pyublas 198 199 a = numpy.random.rand(50000) 200 b = numpy.random.rand(50000) 201 202 dtype = a.dtype 203 204 lin_comb = ElementwiseKernel([ 205 ScalarArg(dtype, "a_fac"), VectorArg(dtype, "a"), 206 ScalarArg(dtype, "b_fac"), VectorArg(dtype, "b"), 207 VectorArg(dtype, "c"), 208 ], 209 "c[i] = a_fac*a[i] + b_fac*b[i]") 210 211 c = numpy.empty_like(a) 212 lin_comb(5, a, 6, b, c) 213 214 import numpy.linalg as la 215 print la.norm(c - (5*a+6*b)) 216

Source Code for Module codepy.elementwise