def __init__(self, backend, shape, iwl, allocator=drv.mem_alloc,

base=None, gpudata=None, strides=None, is_trans=False, order="C"):

dtype = np.dtype(np.int16)

try:

size = 1

for dim in shape:

size *= dim

except TypeError:

assert isinstance(shape, (int, long, np.integer))

size = shape

shape = (shape,)

if isinstance(size, np.integer):

size = np.asscalar(size)

if strides is None:

if order == "F":

strides = _f_contiguous_strides(dtype.itemsize, shape)

elif order == "C":

strides = _c_contiguous_strides(dtype.itemsize, shape)

else:

raise ValueError("invalid order: %s" % order)

else:

strides = tuple(strides)

self.backend = backend

self.base = base

self.shape = shape

self.iwl = iwl

self.strides = strides

self.size = size

self.dtype = dtype

self.nbytes = dtype.itemsize * size

self.allocator = allocator

self.is_trans = is_trans

if gpudata is None:

if size:

self.gpudata = allocator(self.nbytes)

else:

self.gpudata = None

assert base is None

else:

self.gpudata = gpudata

def __str__(self):

return "FlexptArray shape:%s strides:%s is_trans:%s" % (self.shape, self.strides, self.is_trans)

def __repr__(self):

return "FlexptArray"

def __int__(self):

return int(selt.gpudata)

def __len__(self):

"""Return the size of the leading dimension of self."""

if len(self.shape):

return self.shape[0]

else:

return 0

@property

@memoize_method

def flags(self):

return _ArrayFlags(self)

def set(self, ary):

assert ary.dtype == np.int16

assert ary.size == self.size

assert ary.strides == self.strides

assert self.flags.forc

if self.size:

drv.memcpy_htod(self.gpudata, ary)

return self

def get(self, ary=None, astype=None, pagelocked=False):

if ary is None:

if pagelocked:

ary = drv.pagelocked_empty(self.shape, self.dtype)

else:

ary = np.empty(self.shape, self.dtype)

ary = _as_strided(ary, strides=self.strides)

else:

assert ary.size == self.size

assert ary.dtype == self.dtype

assert ary.flags.forc

assert self.flags.forc, "Array in get() must be contiguous"

if self.size:

drv.memcpy_dtoh(ary, self.gpudata)

if astype is not None:

ary = ary.astype(astype) * 2 ** (self.iwl - 15)

return ary

def __getitem__(self, index):

"""

return a sliced view of an array

"""

if not isinstance(index, tuple):

index = (index,)

new_shape = []

new_offset = 0

new_strides = []

seen_ellipsis = False

index_axis = 0

array_axis = 0

while index_axis < len(index):

index_entry = index[index_axis]

if array_axis > len(self.shape):

raise IndexError("too many axes in index")

if isinstance(index_entry, slice):

start, stop, idx_stride = index_entry.indices(

self.shape[array_axis])

array_stride = self.strides[array_axis]

new_shape.append((stop-start)//idx_stride)

new_strides.append(idx_stride*array_stride)

new_offset += array_stride*start

index_axis += 1

array_axis += 1

elif isinstance(index_entry, (int, np.integer)):

array_shape = self.shape[array_axis]

if index_entry < 0:

index_entry += array_shape

if not (0 <= index_entry < array_shape):

raise IndexError(

"subindex in axis %d out of range" % index_axis)

new_offset += self.strides[array_axis]*index_entry

index_axis += 1

array_axis += 1

elif index_entry is Ellipsis:

index_axis += 1

remaining_index_count = len(index) - index_axis

new_array_axis = len(self.shape) - remaining_index_count

if new_array_axis < array_axis:

raise IndexError("invalid use of ellipsis in index")

while array_axis < new_array_axis:

new_shape.append(self.shape[array_axis])

new_strides.append(self.strides[array_axis])

array_axis += 1

if seen_ellipsis:

raise IndexError(

"more than one ellipsis not allowed in index")

seen_ellipsis = True

else:

raise IndexError("invalid subindex in axis %d" % index_axis)

while array_axis < len(self.shape):

new_shape.append(self.shape[array_axis])

new_strides.append(self.strides[array_axis])

array_axis += 1

return self.__class__(

backend = self.backend,

shape = tuple(new_shape),

iwl = self.iwl,

allocator = self.allocator,

base = self,

gpudata = int(self.gpudata)+new_offset,

strides = tuple(new_strides))

def _assign(self, value):

if isinstance(value, (int, float)):

# if we have a c or f contiguous array, then use the speedy driver kernel

if self.flags.forc and float(value) >= 0:

drv.memset_d16(self.gpudata, Flexpt.flex_from_native(value,self.iwl), self.size)

# otherwise use our copy kerel

else:

OpTreeNode.build("copy", value, None, out=self)

elif isinstance(value, FlexptArray):

if self.flags.forc and value.flags.forc and self.iwl == value.iwl:

drv.memcpy_dtod(self.gpudata, value.gpudata, self.nbytes)

else:

OpTreeNode.build("copy", value, None, out=self)

elif isinstance(value, OpTreeNode):

value.execute(out=self)

else:

raise TypeError("Invalid type for assignment: %s" % type(value))

return self

def __setitem__(self, index, value):

self.__getitem__(index)._assign(value)

def fill(self, value):

return self._assign(value)

def copy(self, a):

return self._assign(a)

def reshape(self, *shape):

if isinstance(shape[0], tuple) or isinstance(shape[0], list):

shape = tuple(shape[0])

if shape == self.shape:

return self

size = reduce(lambda x, y: x * y, shape, 1)

if size != self.size:

raise ValueError("total size of new array must be unchanged")

if not self.flags.forc:

raise TypeError("reshaping of non-contigous arrays is not yet supported")

return self.__class__(

backend = self.backend,

shape = shape,

iwl = self.iwl,

allocator = self.allocator,

base = self,

gpudata = self.gpudata,

strides = _c_contiguous_strides(self.dtype.itemsize, shape))

@property

def T(self):

"""

return a transposed view

"""

return self.__class__(

backend = self.backend,

shape = self.shape[::-1],

iwl = self.iwl,

allocator = self.allocator,

base = self,

gpudata = self.gpudata,

strides = self.strides[::-1],

is_trans = not self.is_trans)

def __add__ (self, other): return OpTreeNode.build("add", self, other)

def __sub__ (self, other): return OpTreeNode.build("sub", self, other)

def __mul__ (self, other): return OpTreeNode.build("mul", self, other)

def __div__ (self, other): return OpTreeNode.build("div", self, other)

def __truediv__ (self, other): return OpTreeNode.build("div", self, other)

def __pow__ (self, other): return OpTreeNode.build("pow", self, other)

def __radd__ (self, other): return OpTreeNode.build("add", other, self)

def __rsub__ (self, other): return OpTreeNode.build("sub", other, self)

def __rmul__ (self, other): return OpTreeNode.build("mul", other, self)

def __rdiv__ (self, other): return OpTreeNode.build("div", other, self)

def __rtruediv__ (self, other): return OpTreeNode.build("div", other, self)

def __rpow__ (self, other): return OpTreeNode.build("pow", other, self)

def __eq__ (self, other): return OpTreeNode.build("eq", self, other)

def __ne__ (self, other): return OpTreeNode.build("ne", self, other)

def __lt__ (self, other): return OpTreeNode.build("lt", self, other)

def __le__ (self, other): return OpTreeNode.build("le", self, other)

def __gt__ (self, other): return OpTreeNode.build("gt", self, other)

def __ge__ (self, other): return OpTreeNode.build("ge", self, other)

def __abs__ (self): return OpTreeNode.build("abs", self, None)

def __neg__ (self): return OpTreeNode.build("neg", self, None)

def __iadd__ (self, other): return OpTreeNode.build("add", self, other, out=self)

def __isub__ (self, other): return OpTreeNode.build("sub", self, other, out=self)

def __imul__ (self, other): return OpTreeNode.build("mul", self, other, out=self)

def __idiv__ (self, other): return OpTreeNode.build("div", self, other, out=self)

def __itruediv__ (self, other): return OpTreeNode.build("div", self, other, out=self)

def __init__(self, kernel_set="fgemm_int64_wide32", locks=1024, calc_partials=True, bench=False):

m = re.search( r'wide(\d+)', kernel_set)

if m:

self.width = int(m.group(1))

else:

raise ValueError("Invalid kernel_set")

self.locks = locks

self.module = drv.module_from_file("kernels/" + kernel_set + ".cubin")

self.mode = 0 if calc_partials else 4

self.fgemm = dict()

for op in ("nt", "nn", "tn"):

mod = self.module.get_function(kernel_set + "_" + op)

mod.prepare("PPPIIIIIIHH")

self.fgemm[op] = mod

fprop_conv = self.module.get_function("fprop_conv_float32_K64N64T64")

fprop_conv.prepare("PPPIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII")

self.fgemm["fprop_conv"] = fprop_conv

bprop_conv = self.module.get_function("bprop_conv_float32_CRST64N64T64")

bprop_conv.prepare("PPPPIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII")

self.fgemm["bprop_conv"] = bprop_conv

udpate_conv = self.module.get_function("update_conv_float32_CRST64K64T64")

udpate_conv.prepare("PPPPIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII")

self.fgemm["update_conv"] = udpate_conv

self.gpulock = drv.mem_alloc(locks*4)

drv.memset_d32(self.gpulock, 0, locks)

self.bench = bench

if bench:

self.start = drv.Event()

self.end = drv.Event()

def empty(self, shape, iwl=None, allocator=drv.mem_alloc, order="C"):

"""

allocate the space for a FlexptArray

"""

return FlexptArray(self, shape, iwl, allocator=allocator, order=order)

def array(self, ary, iwl=None, allocator=drv.mem_alloc):

"""

converts a numpy array to a FlexptArray

"""

if iwl is None:

ary, iwl = self.flex_from_native(ary)

elif ary.dtype != np.int16:

ary = self.flex_from_native(ary, iwl)

return FlexptArray(self, ary.shape, iwl, allocator=allocator, strides=ary.strides).set(ary)

def zeros(self, shape, iwl, allocator=drv.mem_alloc, order="C"):

"""

Returns an array of the given shape and dtype filled with 0's.

"""

result = FlexptArray(self, shape, iwl, allocator, order=order)

drv.memset_d16(result.gpudata, 0, result.size)

return result

def ones(self, shape, iwl, allocator=drv.mem_alloc, order="C"):

"""

Returns an array of the given shape and dtype filled with 0's.

"""

result = FlexptArray(self, shape, iwl, allocator, order=order)

drv.memset_d16(result.gpudata, self.flex_from_native(1,iwl), result.size)

return result

def empty_like(self, other_ary):

"""

Returns an array with the same params as another

"""

return FlexptArray(self, other_ary.shape, other_ary.iwl, other_ary.allocator)

def dot(self, A, B, out, iwl=None):

# one dimention must be contiguous

assert min(A.strides) == 2

assert min(B.strides) == 2

assert min(out.strides) == 2

opA = 't' if A.is_trans else 'n' # A.strides[0] < A.strides[1]

opB = 't' if B.is_trans else 'n' # B.strides[0] < B.strides[1]

op = opB + opA

# TODO: use nn and swapp A and B if we need tt op (C is transposed though, I think)

assert op != "tt"

m = A.shape[0]

n = B.shape[1]

k = A.shape[1]

assert m == out.shape[0]

assert n == out.shape[1]

assert k == B.shape[0]

lda = max(A.strides) // 2

ldb = max(B.strides) // 2

ldc = max(out.strides) // 2

gridX = n // self.width + (n % self.width != 0)

gridY = m // self.width + (m % self.width != 0)

if iwl is not None:

out.iwl = iwl

scale = 15 + out.iwl - (A.iwl + B.iwl)

if self.bench:

self.start.record()

self.fgemm[op].prepared_call(

(gridX,gridY,1), (64,1,1),

B.gpudata, A.gpudata, out.gpudata,

ldb, lda, ldc,

n, m, k,

scale, self.mode)

if self.bench:

self.end.record()

self.end.synchronize()

msecs = self.end.time_since(self.start)

gflops = (m * n * k * 2.0) / (msecs * 1000000.0)

print "%s (%5d, %5d, %5d): %9.6f msecs %6.1f gflops" % (op,m,n,k,msecs,gflops)

return (msecs, gflops)

return out

def update_conv(self, I, E, F, padding=(0,0,0), strides=(1,1,1), upscale=(1,1,1), iwl=None, repeat=1):

# I dims: C,D,H,W,N

# F dims: C,T,R,S,K

# E dims: K,M,P,Q,N

assert len(I.shape) == len(F.shape) == len(E.shape), "I, F, and E must be of same dimensions"

dims = len(I.shape) - 2

assert 1 <= dims <= 3, "The image data must be between 1 and 3 dimensions"

N = I.shape[-1]

K = F.shape[-1]

assert N % 8 == 0, "N dim must be multiple of 8"

assert K % 8 == 0, "K dim must be multiple of 8"

assert N == E.shape[-1], "I and E minibatch sizes do not match"

C = I.shape[0]

assert C == F.shape[0], "Channel counts of I and F do not match"

assert K == E.shape[0], "Filter banks and output channels do not match"

if iwl is not None:

F.iwl = iwl

scale = 15 + F.iwl - (I.iwl + E.iwl)

DHW, TRS, MPQ, pad, std, scl = ([],[],[],[],[],[])

# initialize any unspecified leading dimensions

for i in range(3 - dims):

for lst in (DHW, TRS, MPQ, std, scl):

lst.append(1)

pad.append(0)

# fill in the rest

DHW.extend(I.shape[1:dims+1])

TRS.extend(F.shape[1:dims+1])

MPQ.extend(E.shape[1:dims+1])

pad.extend(padding[-dims:])

std.extend(strides[-dims:])

scl.extend(upscale[-dims:])

for i in range(len(DHW)):

if __debug__:

dim = (DHW[i]*scl[i] - TRS[i] + 1 + 2*pad[i]) // std[i]

assert MPQ[i] == dim, "Output dim=%d value=%d does not match calculated value=%d" % (i, MPQ[i], dim)

scl[i] = magic32(MPQ[i]+TRS[i]-pad[i]-2, scl[i])

W = DHW[2]

WN = W*N

HWN = DHW[1]*WN

DHWN = DHW[0]*HWN

S = TRS[2]

RS = TRS[1]*S

RST = TRS[0]*RS

CRST = C*RST

Q = MPQ[2]

P = MPQ[1]

PM = MPQ[0]*P

QPM = Q*PM

magic_RST = magic32(CRST+8, RST)

magic_RS = magic32(RST+32, RS )

magic_S = magic32(RS+32, S )

magic_P = magic32(PM, P )

gridI = CRST // 64 + (CRST % 64 != 0)

gridE = K // 64 + (K % 64 != 0)

kernel_args = flatten([

F.gpudata, I.gpudata, E.gpudata, self.gpulock, scale,

N, K, DHW, WN, HWN, DHWN,

C, CRST, RST, magic_RST, RS, magic_RS, S, magic_S,

pad, std, scl,

Q, P, Q*N, P*Q*N, QPM*N, magic_P])

if self.bench:

self.start.record()

for r in range(repeat):

self.fgemm["update_conv"].prepared_call(

(gridI, gridE, PM), (64,1,1), *kernel_args)

if self.bench:

self.end.record()

self.end.synchronize()

msecs = self.end.time_since(self.start) / repeat

gflops = (QPM * K * N * CRST * 2.0) / (msecs * 1000000.0)

print "update_conv NCK: (%d,%d,%d) DHW:%s TRS:%s MPQ:%s: %9.6f msecs %6.1f gflops" % \

(N, C, K, DHW, TRS, MPQ, msecs, gflops)

return (msecs, gflops)

def fprop_conv(self, I, F, O, padding=(0,0,0), strides=(1,1,1), upscale=(1,1,1), iwl=None, repeat=1):

# I dims: C,D,H,W,N

# F dims: C,T,R,S,K

# O dims: K,M,P,Q,N

assert len(I.shape) == len(F.shape) == len(O.shape), "I, F, and O must be of same dimensions"

dims = len(I.shape) - 2

assert 1 <= dims <= 3, "The image data must be between 1 and 3 dimensions"

N = I.shape[-1]

K = F.shape[-1]

assert N % 8 == 0, "N dim must be multiple of 8"

assert K % 8 == 0, "K dim must be multiple of 8"

assert N == O.shape[-1], "I and O minibatch sizes do not match"

C = I.shape[0]

assert C == F.shape[0], "Channel counts of I and F do not match"

assert K == O.shape[0], "Filter banks and output channels do not match"

if iwl is not None:

O.iwl = iwl

scale = 15 + O.iwl - (I.iwl + F.iwl)

DHW, TRS, MPQ, pad, std, scl = ([],[],[],[],[],[])

# initialize any unspecified leading dimensions

for i in range(3 - dims):

for lst in (DHW, TRS, MPQ, std, scl):

lst.append(1)

pad.append(0)

# fill in the rest

DHW.extend(I.shape[1:dims+1])

TRS.extend(F.shape[1:dims+1])

MPQ.extend(O.shape[1:dims+1])

pad.extend(padding[-dims:])

std.extend(strides[-dims:])

scl.extend(upscale[-dims:])

for i in range(len(DHW)):

if __debug__:

dim = (DHW[i]*scl[i] - TRS[i] + 1 + 2*pad[i]) // std[i]

assert MPQ[i] == dim, "Output dim=%d value=%d does not match calculated value=%d" % (i, MPQ[i], dim)

scl[i] = magic32(MPQ[i]+TRS[i]-pad[i]-2, scl[i])

W = DHW[2]

WN = W*N

HWN = DHW[1]*WN

DHWN = DHW[0]*HWN

S = TRS[2]

RS = TRS[1]*S

RST = TRS[0]*RS

CRST = C*RST

Q = MPQ[2]

PQ = MPQ[1]*Q

PQM = MPQ[0]*PQ

magic_RST = magic32(CRST+8, RST)

magic_RS = magic32(RST+32, RS)

magic_S = magic32(RS+32, S)

magic_PQ = magic32(PQM, PQ)

magic_Q = magic32(PQ, Q)

gridI = N // 64 + (N % 64 != 0)

gridF = K // 64 + (K % 64 != 0)

share = (RST // 32 + (RST % 32 != 0)) * 32 * 4

kernel_args = flatten([

O.gpudata, I.gpudata, F.gpudata, scale,

N, K, DHW, WN, HWN, DHWN,

C, CRST, RST, magic_RST, RS, magic_RS, S, magic_S,

pad, std, scl,

Q, PQ, Q*N, PQ*N, PQM*N, magic_PQ, magic_Q])

if self.bench:

self.start.record()

for r in range(repeat):

self.fgemm["fprop_conv"].prepared_call(

(gridF, gridI, PQM), (64,1,1), *kernel_args, shared_size=share)

if self.bench:

self.end.record()

self.end.synchronize()

msecs = self.end.time_since(self.start) / repeat

gflops = (PQM * K * N * CRST * 2.0) / (msecs * 1000000.0)

print "fprop_conv NCK: (%d,%d,%d) DHW:%s TRS:%s MPQ:%s: %9.6f msecs %6.1f gflops" % \

(N, C, K, DHW, TRS, MPQ, msecs, gflops)

return (msecs, gflops)

def bprop_conv(self, F, E, I, padding=(0,0,0), strides=(1,1,1), upscale=(1,1,1), iwl=None, repeat=1):

# I dims: C,D,H,W,N

# F dims: C,T,R,S,K

# O dims: K,M,P,Q,N

assert len(I.shape) == len(F.shape) == len(E.shape), "I, F, and E must be of same dimensions"

dims = len(I.shape) - 2

assert 1 <= dims <= 3, "The image data must be between 1 and 3 dimensions"

N = I.shape[-1]

K = F.shape[-1]

assert N % 8 == 0, "N dim must be multiple of 8"

assert K % 8 == 0, "K dim must be multiple of 8"

assert N == E.shape[-1], "I and E minibatch sizes do not match"

C = I.shape[0]

assert C == F.shape[0], "Channel counts of I and F do not match"

assert K == E.shape[0], "Filter banks and output channels do not match"

if iwl is not None:

I.iwl = iwl

scale = 15 + I.iwl - (F.iwl + E.iwl)

DHW, TRS, MPQ, pad, std, scl = ([],[],[],[],[],[])

# initialize any unspecified leading dimensions

for i in range(3 - dims):

for lst in (DHW, TRS, MPQ, std, scl):

lst.append(1)

pad.append(0)

# fill in the rest

DHW.extend(I.shape[1:dims+1])

TRS.extend(F.shape[1:dims+1])

MPQ.extend(E.shape[1:dims+1])

pad.extend(padding[-dims:])

std.extend(strides[-dims:])

scl.extend(upscale[-dims:])

for i in range(len(DHW)):

if __debug__:

dim = (DHW[i]*scl[i] - TRS[i] + 1 + 2*pad[i]) // std[i]

assert MPQ[i] == dim, "Output dim=%d value=%d does not match calculated value=%d" % (i, MPQ[i], dim)

scl[i] = magic32(MPQ[i]+TRS[i]-pad[i]-2, scl[i])

W = DHW[2]

WN = W*N

HWN = DHW[1]*WN

DHWN = DHW[0]*HWN

S = TRS[2]

RS = TRS[1]*S

RST = TRS[0]*RS

CRST = C*RST

Q = MPQ[2]

PQ = MPQ[1]*Q

PQM = MPQ[0]*PQ

magic_RST = magic32(CRST+8, RST)

magic_RS = magic32(RST+32, RS)

magic_S = magic32(RS+32, S)

magic_PQ = magic32(PQM, PQ)

magic_Q = magic32(PQ, Q)

lockC = len(bin(128 // RST))-2

gridF = CRST // 64 + (CRST % 64 != 0)

gridE = N // 64 + (N % 64 != 0)

share = (RST // 32 + (RST % 32 != 0)) * 32 * 4

kernel_args = flatten([

I.gpudata, F.gpudata, E.gpudata, self.gpulock, scale,

N, K, DHW, WN, HWN, DHWN,

C, CRST, RST, magic_RST, RS, magic_RS, S, magic_S,

pad, std, scl,

Q, PQ, Q*N, PQ*N, PQM*N, magic_PQ, magic_Q, PQM*N*8*2, magic_RST[1]+lockC])

if self.bench:

self.start.record()

for r in range(repeat):

self.fgemm["bprop_conv"].prepared_call(

(gridF, gridE, PQM), (64,1,1), *kernel_args, shared_size=share)

if self.bench:

self.end.record()

self.end.synchronize()

msecs = self.end.time_since(self.start) / repeat

gflops = (PQM * K * N * CRST * 2.0) / (msecs * 1000000.0)

print "bprop_conv NCK: (%d,%d,%d) DHW:%s TRS:%s MPQ:%s: %9.6f msecs %6.1f gflops" % \

(N, C, K, DHW, TRS, MPQ, msecs, gflops)

return (msecs, gflops)

def clip(self, a, a_min, a_max, out=None):

return OpTreeNode.build("min", OpTreeNode.build("max", a, a_min), a_max, out=out)

def maximum(self, a, b, out=None): return OpTreeNode.build("max", a, b, out=out)

def minimum(self, a, b, out=None): return OpTreeNode.build("min", a, b, out=out)

def add (self, a, b, out=None, iwl=None): return OpTreeNode.build("add", a, b, out=out, iwl=iwl)

def subtract (self, a, b, out=None, iwl=None): return OpTreeNode.build("sub", a, b, out=out, iwl=iwl)

def multiply (self, a, b, out=None, iwl=None): return OpTreeNode.build("mul", a, b, out=out, iwl=iwl)

def divide (self, a, b, out=None, iwl=None): return OpTreeNode.build("div", a, b, out=out, iwl=iwl)

def true_divide (self, a, b, out=None, iwl=None): return OpTreeNode.build("div", a, b, out=out, iwl=iwl)

def power (self, a, b, out=None, iwl=None): return OpTreeNode.build("pow", a, b, out=out, iwl=iwl)

def reciprocal (self, a, out=None, iwl=None): return OpTreeNode.build("div", 1, a, out=out, iwl=iwl)

def negative (self, a, out=None): return OpTreeNode.build("neg", a, None, out=out)

def absolute (self, a, out=None): return OpTreeNode.build("abs", a, None, out=out)

def fabs (self, a, out=None): return OpTreeNode.build("abs", a, None, out=out)

def sqrt (self, a, out=None, iwl=None): return OpTreeNode.build("sqrt", a, None, out=out, iwl=iwl)

def sqare(self, a, out=None, iwl=None): return OpTreeNode.build("sqr", a, None, out=out, iwl=iwl)

def exp (self, a, out=None, iwl=None): return OpTreeNode.build("exp", a, None, out=out, iwl=iwl)

def exp2 (self, a, out=None, iwl=None): return OpTreeNode.build("exp2", a, None, out=out, iwl=iwl)

def log (self, a, out=None, iwl=None): return OpTreeNode.build("log", a, None, out=out, iwl=iwl)

def log2 (self, a, out=None, iwl=None): return OpTreeNode.build("log2", a, None, out=out, iwl=iwl)

def sig (self, a, out=None, iwl=None): return OpTreeNode.build("sig", a, None, out=out, iwl=iwl)

def sig2 (self, a, out=None, iwl=None): return OpTreeNode.build("sig2", a, None, out=out, iwl=iwl)

def tanh (self, a, out=None, iwl=None): return OpTreeNode.build("tanh", a, None, out=out, iwl=iwl)

def tanh2(self, a, out=None, iwl=None): return OpTreeNode.build("tanh2",a, None, out=out, iwl=iwl)

def equal (self, a, b, out=None): return OpTreeNode.build("eq", a, b, out=out)

def not_equal (self, a, b, out=None): return OpTreeNode.build("ne", a, b, out=out)

def less (self, a, b, out=None): return OpTreeNode.build("lt", a, b, out=out)

def less_equal (self, a, b, out=None): return OpTreeNode.build("le", a, b, out=out)

def greater (self, a, b, out=None): return OpTreeNode.build("gt", a, b, out=out)

def greater_equal (self, a, b, out=None): return OpTreeNode.build("ge", a, b, out=out)

#TODO: reshape to more efficient dimensions if needed

def sum(self, a, out=None, iwl=None, axis=None, partial=None):

return OpTreeNode.build("sum", a, None, out=out, iwl=iwl, axis=axis, partial=partial)

def max(self, a, out=None, iwl=None, axis=None, partial=None):

return OpTreeNode.build("max", a, None, out=out, iwl=iwl, axis=axis, partial=partial)

def min(self, a, out=None, iwl=None, axis=None, partial=None):

return OpTreeNode.build("min", a, None, out=out, iwl=iwl, axis=axis, partial=partial)

@staticmethod

def flex_from_native(value, iwl=None):

if type(value) not in (int, float, np.ndarray):

raise TypeError("Unsupported type: %s" % type(value))

# find an appropriate iwl and convert

if iwl is None:

# find the absolute max and special case all zeros

if type(value) is np.ndarray:

max = np.max(np.absolute(value))

if max == 0:

return (value.astype(np.int16),0)

else:

if value == 0:

return (0,0)

max = value

# convert the maximum to float and extract the exponent

iwl = ((struct.unpack('I',struct.pack('f',float(max)))[0] & 0x7f800000) >> 23)-126

if type(value) is np.ndarray:

return ((value * 2.0 ** (15 - iwl)).astype(np.int16), iwl)

return (int(value * 2 ** (15 - iwl)), iwl)

if type(iwl) is not int:

raise TypeError("iwl must be an int")

# convert from int

if type(value) is int:

if iwl < 1 or iwl > 15:

raise ValueError("value(%d) does not fit in flex with requested iwl(%d)" % (value, iwl))

return value << (15 - iwl)

#convert from float

if type(value) is float:

return int(value * 2 ** (15 - iwl))

#ndarray

return (value * 2.0 ** (15 - iwl)).astype(np.int16)

@staticmethod

def native_from_flex(flex, iwl):

#TODO: this needs more work for other types

def __new__(cls, *args):

return tuple.__new__(cls, args)

@staticmethod

def build(op, a, b, out=None, **kwargs):

for arg in (a,b):

if not isinstance(arg, (int, float, FlexptArray, OpTreeNode, type(None))):

return NotImplemented

op_dict = { "op" : op }

op_dict.update(kwargs)

node = OpTreeNode(op_dict, a, b)

# delay execution until assignment

if out is None:

return node

# passing in an out value counts as assignment

return node.execute(out=out)

def execute(self, out):

stack = self.traverse(list())

return flexpt_ew.call_compound_ew_kernel(out, *stack)

# post order walk op tree and produce postfix stack

def traverse(self, stack):

# Left

if type(self[1]) is OpTreeNode:

self[1].traverse(stack)

else:

stack.append(self[1])

# Right

if type(self[2]) is OpTreeNode:

self[2].traverse(stack)

elif self[2] is not None:

stack.append(self[2])

stack.append(self[0])

return stack

def __add__ (self, other): return self.build("add", self, other)

def __sub__ (self, other): return self.build("sub", self, other)

def __mul__ (self, other): return self.build("mul", self, other)

def __div__ (self, other): return self.build("div", self, other)

def __truediv__ (self, other): return self.build("div", self, other)

def __pow__ (self, other): return self.build("pow", self, other)

def __radd__ (self, other): return self.build("add", other, self)

def __rsub__ (self, other): return self.build("sub", other, self)

def __rmul__ (self, other): return self.build("mul", other, self)

def __rdiv__ (self, other): return self.build("div", other, self)

def __rtruediv__ (self, other): return self.build("div", other, self)

def __rpow__ (self, other): return self.build("pow", other, self)

def __eq__ (self, other): return self.build("eq", self, other)

def __ne__ (self, other): return self.build("ne", self, other)

def __lt__ (self, other): return self.build("lt", self, other)

def __le__ (self, other): return self.build("le", self, other)

def __gt__ (self, other): return self.build("gt", self, other)

def __ge__ (self, other): return self.build("ge", self, other)

def __abs__ (self): return self.build("abs", self, None)

nc = ((nmax + 1)//d)*d - 1

nbits = len(bin(nmax)) - 2

for p in range(0, 2*nbits + 1):

if 2**p > nc*(d - 1 - (2**p - 1)%d):

m = (2**p + d - 1 - (2**p - 1)%d)//d

return (m, p)

# 3 is a special case that only ends up in the high bits

# if the maxn is 0xffffffff

# we can't use 0xffffffff for all cases as some return a 33 bit

# magic number

maxn = 0xffffffff if d == 3 else 0x7fffffff

magic, shift = magic32(maxn, d)

if magic != 1:

shift -= 32