softmin = 1

sparsemin = 2

smoothmin = 3

n_features = r.shape[0]

MAX_SIZE = 16

local_x = cuda.local.array(MAX_SIZE, dtype=types.float32)

local_u = cuda.local.array(MAX_SIZE, dtype=types.float32)

local_cssv = cuda.local.array(MAX_SIZE, dtype=types.float32)

allInf = True

for i in range(n_features):

allInf *= r[i] > 1e10

if allInf:

val[0] = math.inf

for i in range(n_features):

grads[i] = 1/n_features

return

for i in range(n_features):

if r[i]>1e20:

local_x[i] = 1e20

else:

local_x[i] = r[i]

for i in range(n_features):

local_x[i] = -local_x[i] / gamma

for i in range(n_features):

local_u[i] = local_x[i]

for i in range(1, n_features):

key = local_u[i]

j = i - 1

while j >= 0 and local_u[j] < key:

local_u[j + 1] = local_u[j]

j -= 1

local_u[j + 1] = key

local_cssv[0] = local_u[0] - 1.0

for i in range(1, n_features):

local_cssv[i] = local_cssv[i - 1] + local_u[i]

cMax = local_cssv[0]

rho = 1

for i in range(1, n_features):

if local_u[i] - local_cssv[i]/(i+1) > 0:

rho = i+1

cMax = local_cssv[i]

else:

break

theta = cMax / rho

val[0] = -gamma/2

for i in range(n_features):

grads[i] = max(local_x[i]-theta,0)

n_steps = d_dir.shape[0]

soft_result = cuda.local.array((1,), types.float32)

soft_grad = cuda.local.array((16,), types.float32)

for s in range(n_steps):

d_dir[s] = min(d_dir[s], 1e20)

softminX(d_dir, gamma, soft_result, soft_grad)

for s in range(n_steps):

val[0] += d_dir[s]*soft_grad[s]

for s in range(n_steps):

n_steps = d_dir.shape[0]

allInf = True

for i in range(n_steps):

allInf *= math.isinf(d_dir[i])

if allInf:

val[0] = math.inf

for i in range(n_steps):

grads[i] = 1/n_steps

else:

min_d = math.inf

for i in range(n_steps):

min_d = min(min_d, d_dir[i])

sum_exp = 0

for i in range(n_steps):

sum_exp += math.exp(-(d_dir[i]-min_d)/gamma)

val[0] = -gamma * (-min_d/gamma + math.log(sum_exp))

for i in range(n_steps):

n_steps = d_dir.shape[0]

allInf = True

for i in range(n_steps):

allInf *= math.isinf(d_dir[i])

if allInf:

val[0] = math.inf

for i in range(n_steps):

grads[i] = 1/n_steps

else:

min_index = 0

min_val = d_dir[0]

for i in range(n_steps):

if d_dir[i] < min_val:

min_val = d_dir[i]

min_index = i

val[0] = min_val

min_func, list_N, list_M, steps, sub_start_m, sub_end_m):

# Each block processes one pair of examples

b = cuda.blockIdx.x

# We have as many threads as seq_len, because the most number of threads we need

# is equal to the number of elements on the largest anti-diagonal

tid = cuda.threadIdx.x

# Compute I, J, the indices from [0, seq_len)

# The row index is always the same as tid

I = tid

inv_gamma = 1.0 / gamma

pad_N = int(max(steps[:,0]))

pad_M = int(max(steps[:,1]))

# Go over each anti-diagonal. Only process threads that fall on the current on the anti-diagonal

for p in range(n_passes):

# The index is actually 'p - tid' but need to force it in-bounds

J = max(0, min(p - tid, max_j - 1))

# For simplicity, we define i, j which start from 1 (offset from I, J)

i = I + pad_N

j = J + pad_M

# Only compute if element[i, j] is on the current anti-diagonal, and also is within bounds

if (I + J == p) and (I < max_i and J < max_j) and (I<list_N[b]) and (J<list_M[b]):

# don't compute first elements

if not ((i==pad_N) and (j<pad_M+sub_start_m)):

# Don't compute if outside bandwidth

if not (abs(i - j) > bandwidth > 0):

for s in range(steps.shape[0]):

d_dir[b,tid,s] = W_[b,i,j,s]*C_[b,i,j] + D_[b, i-steps[s,0], j-steps[s,1]]

if min_func == minFunc.softmin:

softminX(d_dir[b,tid], gamma, D_[b,i,j:j+1], K_[b,i,j])

elif min_func == minFunc.hardmin:

hardminX(d_dir[b,tid], gamma, D_[b,i,j:j+1], K_[b,i,j])

elif min_func==minFunc.smoothmin:

smoothminX(d_dir[b,tid], gamma, D_[b,i,j:j+1], K_[b,i,j])

elif min_func==minFunc.sparsemin:

sparseminX(d_dir[b,tid], gamma, D_[b,i,j:j+1], K_[b,i,j])

for s in range(steps.shape[0]):

G_[b,i,j] += K_[b,i,j,s]*W_[b,i,j,s]

# Wait for other threads in this block

cuda.syncthreads()

# subsequence-DTW ending: get the final cost value and initialize the backward recursion;

# we only need one thread for this (here, we use the thread from I==1)

if I == 0:

# no need to for calculations of no-subsequence ending

if sub_end_m == 1:

D_[b,-1,-1] = D_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-1]

E_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-1] = 1

else:

if min_func == minFunc.softmin:

softminX(D_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M],

gamma,

D_[b,-1,-1:],

E_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M])

elif min_func == minFunc.hardmin:

hardminX(D_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M],

gamma,

D_[b,-1,-1:],

E_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M])

elif min_func == minFunc.smoothmin:

smoothminX(D_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M],

gamma,

D_[b,-1,-1:],

E_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M])

elif min_func == minFunc.sparsemin:

sparseminX(D_[b,list_N[b]+pad_N-1, list_M[b]+pad_M-sub_end_m : list_M[b]+pad_M],

gamma,

D_[b,-1,-1:],

steps, sub_start_m, sub_end_m):

b = cuda.blockIdx.x

tid = cuda.threadIdx.x

pad_N = int(max(steps[:,0]))

pad_M = int(max(steps[:,1]))

# Indexing logic is the same as above, however, the anti-diagonal needs to

# progress backwards

I = tid

for p in range(n_passes):

# Reverse the order to make the loop go backward

rev_p = n_passes - p - 1

# convert tid to I, J, then i, j

J = max(0, min(rev_p - tid, max_j - 1))

i = I + pad_N

j = J + pad_M

# Only compute if element[i, j] is on the current anti-diagonal, and also is within bounds

if (I + J == rev_p) and (I < max_i and J < max_j) and (I<list_N_tensor[b]) and (J<list_M_tensor[b]):

if not ( (I==list_N_tensor[b]-1) and (J>=list_M_tensor[b]-sub_end_m)):

# Don't compute if outside bandwidth

if not (abs(i - j) > bandwidth > 0):

for s in range(steps.shape[0]):

E_[b,i,j] += E_[b, i+steps[s,0], j+steps[s,1]] * K_[b, i+steps[s,0], j+steps[s,1], s]

# Wait for other threads in this block

"""

e_matrix = None

h_matrix = None

d_matrix = None

c_matrix = None

w_matrix = None

@staticmethod

def forward(ctx, C, gamma, bandwidth, weights, min_func, list_N, list_M, steps, sub_start_m, sub_end_m, W_overwrite, verbose):

ctx.verbose=verbose

dev = C.device

dtype = C.dtype

start = timer()

ctx.list_M_tensor = list_M

ctx.list_N_tensor = list_N

ctx.steps = torch.tensor(steps, device=dev, dtype=torch.int16)

ctx.sub_start_m = sub_start_m

ctx.sub_end_m = sub_end_m

pad_N = torch.max(ctx.steps[:,0]).item()

pad_M = torch.max(ctx.steps[:,1]).item()

ctx.pad_N = pad_N

ctx.pad_M = pad_M

B = C.shape[0]

N = C.shape[1]

M = C.shape[2]

threads_per_block = max(N, M)

n_passes = 2 * threads_per_block - 1

# Prepare the output arrays

D_ = torch.zeros((B, N+pad_N*2, M+pad_M*2), device=dev)

C_ = torch.zeros((B, N+pad_N*2, M+pad_M*2), device=dev)

E_ = torch.zeros((B, N+pad_N*2, M+pad_M*2), device=dev)

G_ = torch.zeros((B, N+pad_N*2, M+pad_M*2), device=dev)

W_ = torch.ones((B, N+pad_N*2, M+pad_M*2, ctx.steps.shape[0]), device=dev)

K_ = torch.zeros((B, N+pad_N*2, M+pad_M*2, ctx.steps.shape[0]), device=dev)

d_dir = torch.zeros((B, threads_per_block, ctx.steps.shape[0]),device=dev)

C_[:,pad_N:-pad_N, pad_M:-pad_M] = C

D_[:,:,:pad_M] = math.inf

D_[:,:pad_N,:] = math.inf

# initialize the first values of D that are not updated in the recursion (the allowed start elements in C)

D_[:,pad_N, pad_M:pad_M+ctx.sub_start_m] = C_[:,pad_N, pad_M:pad_M+ctx.sub_start_m]

D_[:,pad_N,pad_M] = C_[:,pad_N,pad_M]

G_[:,pad_N,pad_M:pad_M+ctx.sub_start_m] = 1.

if W_overwrite is None:

for iw, w in enumerate(weights):

W_[:,pad_N:-pad_N,pad_M:-pad_M,iw] = w

else:

W_[:,pad_N:-pad_N, pad_M:-pad_M] = W_overwrite

if ctx.verbose: print("fwd init : %.4fms"%((timer()-start)*1e3))

start = timer()

# Run the CUDA kernel.

# Set CUDA's grid size to be equal to the batch size (every CUDA block processes one sample pair)

# Set the CUDA block size to be equal to the length of the longer sequence (equal to the size of the largest diagonal)

compute_dDTW_cuda[B, threads_per_block](gamma.item(), bandwidth.item(), N, M, n_passes,

cuda.as_cuda_array(W_),

cuda.as_cuda_array(C_),

cuda.as_cuda_array(D_),

cuda.as_cuda_array(E_),

cuda.as_cuda_array(G_),

cuda.as_cuda_array(K_),

cuda.as_cuda_array(d_dir),

min_func,

ctx.list_N_tensor,

ctx.list_M_tensor,

ctx.steps,

ctx.sub_start_m,

ctx.sub_end_m)

if ctx.verbose: print("fwd loop: %.4fms"%((timer()-start)*1e3))

start = timer()

ctx.save_for_backward(bandwidth, E_.clone(), G_.clone(), K_.clone())

_dDTWCUDA.c_matrix = C_

_dDTWCUDA.d_matrix = D_

_dDTWCUDA.W_matrix = W_

if ctx.verbose: print("fwd save: %.4fms"%((timer()-start)*1e3))

return D_[:,-1,-1]

@staticmethod

def backward(ctx, grad_output):

start = timer()

dev = grad_output.device

dtype = grad_output.dtype

bandwidth, E_, G_, K_= ctx.saved_tensors

pad_N = ctx.pad_N

pad_M = ctx.pad_M

B = G_.shape[0]

N = G_.shape[1] - 2*pad_N

M = G_.shape[2] - 2*pad_M

threads_per_block = max(N, M)

n_passes = 2 * threads_per_block - 1

if ctx.verbose: print("bwd init: %.4fms"%((timer()-start)*1e3))

start = timer()

# Grid and block sizes are set same as done above for the forward() call

compute_dDTW_backward_cuda[B, threads_per_block](bandwidth.item(), N, M, n_passes,

cuda.as_cuda_array(E_),

cuda.as_cuda_array(G_),

cuda.as_cuda_array(K_),

ctx.list_N_tensor, ctx.list_M_tensor, ctx.steps,

ctx.sub_start_m, ctx.sub_end_m)

if ctx.verbose: print("bwd loop: %.4fms"%((timer()-start)*1e3))

start = timer()

H = E_[:, pad_N:N+pad_N, pad_M:M+pad_M] * G_[:,pad_N:N+pad_N, pad_M:M+pad_M]

_dDTWCUDA.h_matrix = H

E = E_[:, pad_N:N + pad_N, pad_M:M + pad_M]

_dDTWCUDA.e_matrix = E

if ctx.verbose: print("bwd save: %.4fms"%((timer()-start)*1e3))

"""

def __init__(self, use_cuda, gamma=1.0, bandwidth=None, dist_func=None, min_func=None,

weights=[1,1,1], steps=[[1,0],[0,1],[1,1]], sub_start_m=1, sub_end_m=1, verbose=False, normalization="N", squeeze_1=True):

"""

self.verbose = verbose

self.squeeze_1 = squeeze_1

assert normalization in ["M", "N", "MN", "ctc", "none"]

self.normalization = normalization

self.steps = steps

self.sub_start_m = sub_start_m

self.sub_end_m = sub_end_m

super(dDTW, self).__init__()

self.gamma = torch.tensor([gamma], device="cuda" if use_cuda else "cpu", dtype=torch.float16)

if bandwidth is None:

self.bandwidth = torch.tensor([0], device="cuda" if use_cuda else "cpu", dtype=torch.float16)

else:

self.bandwidth = torch.tensor([bandwidth], device="cuda" if use_cuda else "cpu", dtype=torch.float16)

self.use_cuda = use_cuda

if not use_cuda:

print("Only CUDA mode implemented")

return None

self.dtw_class = None

self.weights = weights

if min_func is None:

self.min_func = minFunc.softmin

else:

self.min_func = minFunc[min_func]

# Set the distance function

if dist_func is not None:

self.dist_func = dist_func

else:

self.dist_func = dDTW._euclidean_dist_func

def _get_func_dtw(self, x, y):

"""

if len(x.shape) == 3:

bx, lx, dx = x.shape

by, ly, dy = y.shape

# Make sure the dimensions match

assert bx == by # Equal batch sizes

assert dx == dy # Equal feature dimensions

elif len(x.shape) == 4:

bx, lx, dx, cx = x.shape

by, ly, dy, cy = y.shape

# Make sure the dimensions match

assert bx == by # Equal batch sizes

assert dx == dy # Equal feature dimensions

assert cx == cy # Equal channel dimensions

else:

print("dimensionality of x and/or y not supported")

use_cuda = self.use_cuda

if use_cuda and (lx > 1024 or ly > 1024): # We should be able to spawn enough threads in CUDA

print("SoftDTW: Cannot use CUDA because the sequence length > 1024 (the maximum block size supported by CUDA)")

return None

# Finally, return the correct function

self.dtw_class = _dDTWCUDA

return _dDTWCUDA.apply

@staticmethod

def _euclidean_dist_func(x, y):

"""

n = x.size(1)

m = y.size(1)

d = x.size(2)

x = x.unsqueeze(2).expand(-1, n, m, d)

y = y.unsqueeze(1).expand(-1, n, m, d)

return torch.pow(x - y, 2).sum(3)

def forward(self, X, Y, list_N=None, list_M=None, W_overwrite=None):

"""

start=timer()

# Format input to the right shape

if self.squeeze_1:

X = torch.squeeze(X, 1)

Y = torch.squeeze(Y, 1)

num_frames = X.shape[1]

num_targets = Y.shape[1]

#print(X.shape)

if list_N is None:

list_N=torch.tensor([num_frames for _ in range(X.shape[0])], device="cuda" if self.use_cuda else "cpu", dtype=torch.int16)

if list_M is None:

list_M=torch.tensor([num_targets for _ in range(Y.shape[0])], device="cuda" if self.use_cuda else "cpu", dtype=torch.int16)

# Check the inputs and get the correct implementation

func_dtw = self._get_func_dtw(X, Y)

if self.verbose: print("loss init: %.4fms"%((timer()-start)*1e3))

start = timer()

C_xy = self.dist_func(X, Y)

if self.verbose: print("cost matrix: %.4fms"%((timer()-start)*1e3))

start = timer()

dtw_cost = func_dtw(C_xy, self.gamma, self.bandwidth, self.weights, self.min_func, list_N, list_M, self.steps, self.sub_start_m, self.sub_end_m, W_overwrite, self.verbose)

if self.verbose: print("dtw cost: %.4fms"%((timer()-start)*1e3))

start = timer()

if self.normalization == "N":

dtw_cost_mean = torch.mean(dtw_cost/list_N)

elif self.normalization == "M":

dtw_cost_mean = torch.mean(dtw_cost/list_M)

elif self.normalization == "MN":

dtw_cost_mean = torch.mean(dtw_cost/list_N/list_M)

elif self.normalization == "ctc":

dtw_cost_mean = torch.mean(dtw_cost/( (list_M-1)/2))

else:

dtw_cost_mean = torch.mean(dtw_cost)

if self.verbose: print("dtw cost mean: %.4fms"%((timer()-start)*1e3))

return dtw_cost_mean