// cudamatrix/cu-kernels.cu

// Copyright 2009-2012 Karel Vesely

// 2013 Ehsan Variani

// 2013 Johns Hopkins University (author: Daniel Povey)

// 2013 Hainan Xu

// 2013 Xiaohui Zhang

// 2013-2015 Guoguo Chen

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED

// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,

// MERCHANTABLITY OR NON-INFRINGEMENT.

// See the Apache 2 License for the specific language governing permissions and

// limitations under the License.

// In this file is the CUDA code of the CUDA kernels, plus the ANSI-C wrappers

#include <cfloat>

#include <limits>

#include <math_constants.h>

#include "cudamatrix/cu-kernels-ansi.h"

/***********************************************************************

* Generic __device__ functions

*/

template<typename Real>

__device__

static Real _sum_reduce(Real buffer[]) {

// Total number of active threads

int32_cuda nTotalThreads = blockDim.x;

__syncthreads();

// perform tree-based reduction (sum)

while(nTotalThreads > 1) {

int32_cuda halfPoint = ((1+nTotalThreads) >> 1); // divide by two

// only the first half of the threads will be active.

if (threadIdx.x >= halfPoint) { // was <

// Get the shared value stored by another thread

Real temp = 0.0;

if(threadIdx.x < nTotalThreads) { // was +halfPoint

temp = buffer[threadIdx.x]; // was +halfPoint

}

buffer[threadIdx.x - halfPoint] += temp;

}

__syncthreads();

nTotalThreads = ((1+nTotalThreads) >> 1); // divide by two.

}

// the result

return buffer[0];

}

/***********************************************************************

* CUDA kernels

* the functions are templated to have the float/double operations

*/

/*

* CuMatrix

*/

template<typename Real>

__global__

static void _copy_low_upp(Real* A, MatrixDim dimA) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

int j = blockIdx.y * blockDim.y + threadIdx.y;

if (i <= j || i >= dimA.rows) return;

int index_1 = i * dimA.stride + j;

int index_2 = j * dimA.stride + i;

A[index_2] = A[index_1];

}

template<typename Real>

__global__

static void _copy_upp_low(Real* A, MatrixDim dimA) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

int j = blockIdx.y * blockDim.y + threadIdx.y;

if (j <= i || j >= dimA.rows) return;

int index_1 = i * dimA.stride + j;

int index_2 = j * dimA.stride + i;

A[index_2] = A[index_1];

}

// mat += diag(vec) * mat2.

template<typename Real>

__global__

static void _add_diag_vec_mat(Real alpha, Real *mat, MatrixDim mat_dim,

const Real *vec, const Real *mat2, int mat2_row_stride,

int mat2_col_stride, Real beta) {

int i = blockIdx.x * blockDim.x + threadIdx.x; // column index

int j = blockIdx.y * blockDim.y + threadIdx.y; // row index

int index = j * mat_dim.stride + i,

index2 = j * mat2_row_stride + i * mat2_col_stride;

if (i < mat_dim.cols && j < mat_dim.rows) {

mat[index] = alpha * vec[j] * mat2[index2] + beta * mat[index];

}

}

template<typename Real, typename OtherReal>

__global__

static void _copy_from_tp(Real* A, const OtherReal* B, MatrixDim dmat) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x; // col index

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y; // row index

if (i < dmat.cols && j < dmat.rows) {

int32_cuda index_B = (j * (j+1) / 2) + i;

int32_cuda index_A = j * dmat.stride + i;

if (i <= j) {

A[index_A] = B[index_B];

} else {

A[index_A] = 0.0;

}

}

}

template<typename Real, typename OtherReal>

__global__

static void _copy_from_tp_trans(Real* A, const OtherReal* B, MatrixDim dmat) {

// we interpret these indexes oppositely from normal, but it doesn't

// matter as it's invoked in a symmetric way.

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

// transpose the indices used to index the source TpMatrix.

if (i < dmat.rows && j < dmat.cols) {

int32_cuda index_B = (j * (j+1) / 2) + i;

int32_cuda index_A = i * dmat.stride + j;

if (i <= j) {

A[index_A] = B[index_B];

} else {

A[index_A] = 0.0;

}

}

}

template<typename Real, typename OtherReal>

__global__

static void _copy_from_mat(Real* mat_out, const OtherReal* mat_in, MatrixDim d_out, MatrixDim d_in) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x; // col-index

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y; // row-index.

int32_cuda index_out = i + j * d_out.stride;

int32_cuda index_in = i + j * d_in.stride;

if (i < d_out.cols && j < d_out.rows)

mat_out[index_out] = static_cast<Real>(mat_in[index_in]);

}

template<int TileDim, typename Real, typename OtherReal>

__global__

static void _copy_from_mat_trans(Real* mat_out, const OtherReal* mat_in,

MatrixDim d_out, MatrixDim d_in) {

// Use shared meme to achieve both coalesced memory reading and writing

// '+1' to avoid bank conflict when reading sbuf

__shared__ Real sbuf[TileDim][TileDim + 1];

const int32_cuda i_in = blockIdx.y * TileDim + threadIdx.y; // row-index

const int32_cuda j_in = blockIdx.x * TileDim + threadIdx.x; // col-index

const int32_cuda tile_stride_in = CU1DBLOCK / TileDim * d_in.stride;

int32_cuda index_in = i_in * d_in.stride + j_in;

# pragma unroll

for (int i = 0; i < TileDim; i += CU1DBLOCK / TileDim) {

if (i_in + i < d_in.rows && j_in < d_in.cols) {

sbuf[threadIdx.y + i][threadIdx.x] = static_cast<Real>(mat_in[index_in]);

}

index_in += tile_stride_in;

}

__syncthreads();

// Grid is transposed, but block is not yet.

// Warp (blockDim.x) is always along the row-dim.

const int32_cuda i_out = blockIdx.x * TileDim + threadIdx.y;

const int32_cuda j_out = blockIdx.y * TileDim + threadIdx.x;

const int32_cuda tile_stride_out = CU1DBLOCK / TileDim * d_out.stride;

int32_cuda index_out = i_out * d_out.stride + j_out;

# pragma unroll

for (int i = 0; i < TileDim; i += CU1DBLOCK / TileDim) {

if (i_out + i < d_out.rows && j_out < d_out.cols) {

// block is tranposed when reading sbuf

mat_out[index_out] = sbuf[threadIdx.x][threadIdx.y + i];

}

index_out += tile_stride_out;

}

}

template<typename Real, typename OtherReal>

__global__

static void _copy_from_smat(Real* mat_out, const MatrixElement<OtherReal>* smat_in, MatrixDim d_out, MatrixIndexT_cuda d_in) {

int smat_index = blockIdx.x * blockDim.x + threadIdx.x;

if (smat_index >= d_in) return;

int data_index = smat_in[smat_index].row * d_out.stride + smat_in[smat_index].column;

mat_out[data_index] = smat_in[smat_index].weight;

}

template<typename Real, typename OtherReal>

__global__

static void _copy_from_smat_trans(Real* mat_out, const MatrixElement<OtherReal>* smat_in, MatrixDim d_out, MatrixIndexT_cuda d_in) {

int smat_index = blockIdx.x * blockDim.x + threadIdx.x;

if (smat_index >= d_in) return;

int data_index = smat_in[smat_index].column * d_out.stride + smat_in[smat_index].row;

mat_out[data_index] = smat_in[smat_index].weight;

}

template<typename Real>

__global__

static void _trace_mat_smat_trans(const Real* mat_in, const MatrixElement<Real>* smat_in, MatrixDim mat_d_in, MatrixIndexT_cuda smat_d_in, Real* trace_vec_out) {

int smat_index = blockIdx.x * blockDim.x + threadIdx.x;

if (smat_index >= smat_d_in) return;

int mat_index = smat_in[smat_index].row * mat_d_in.stride + smat_in[smat_index].column;

trace_vec_out[smat_index] = mat_in[mat_index] * smat_in[smat_index].weight;

}

template<typename Real>

__global__

static void _trace_mat_smat(const Real* mat_in, const MatrixElement<Real>* smat_in, MatrixDim mat_d_in, MatrixIndexT_cuda smat_d_in, Real* trace_vec_out) {

int smat_index = blockIdx.x * blockDim.x + threadIdx.x;

if (smat_index >= smat_d_in) return;

int mat_index = smat_in[smat_index].column * mat_d_in.stride + smat_in[smat_index].row;

trace_vec_out[smat_index] = mat_in[mat_index] * smat_in[smat_index].weight;

}

template<typename Real>

__global__

static void _apply_exp(Real* mat, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j * d.stride;

if (i < d.cols && j < d.rows) {

mat[index] = exp(mat[index]);

}

}

template<typename Real>

__global__

static void _scale_diag_packed(Real* mat, Real value, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda index = ((i+1)*(i+2)/2) - 1;

if ( i < dim ) {

mat[index] = value * mat[index];

}

}

template<typename Real>

__global__

static void _set_diag(Real* mat, Real value, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda index = i + i*d.stride;

if ( i < d.rows && i < d.cols) {

mat[index] = value;

}

}

template<typename Real>

__global__

static void _set_diag_packed(Real* mat, Real value, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda index = ((i+1)*(i+2)/2) - 1;

if ( i < dim ) {

mat[index] = value;

}

}

template<typename Real>

__global__

static void _add_diag_packed(Real* mat, Real value, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda index = ((i+1)*(i+2)/2) - 1;

if ( i < dim ) {

mat[index] = mat[index] + value;

}

}

template<typename Real>

__global__

static void _set_const(Real* mat, Real value, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x; // column

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y; // row

int32_cuda index = i + j * d.stride;

if (i < d.cols && j < d.rows)

mat[index] = value;

}

template<typename Real>

__global__

static void _set_zero_above_diag(Real* mat, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j * d.stride;

if (i < d.cols && j < i)

mat[index] = 0.0;

}

template<typename Real>

__global__

static void _add(Real* mat, Real value, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

mat[index] = mat[index] + value;

}

template<typename Real>

__global__

static void _scale(Real* mat, Real value, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

mat[index] = mat[index] * value;

}

template<typename Real>

__global__

static void _apply_log(Real* mat, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

mat[index] = log(mat[index]);

}

template<typename Real>

__global__

static void _mul_elements(Real* mat, const Real* A, MatrixDim dst_d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda dst_index = i + j*dst_d.stride, src_index = i + j*src_stride;

if (i < dst_d.cols && j < dst_d.rows)

mat[dst_index] = mat[dst_index] * A[src_index];

}

template<typename Real>

__global__

static void _div_elements(Real* mat, const Real* A, MatrixDim dst_d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda dst_index = i + j*dst_d.stride, src_index = i + j*src_stride;

if (i < dst_d.cols && j < dst_d.rows)

mat[dst_index] = mat[dst_index] / A[src_index];

}

template<typename Real>

__global__

static void _max(Real* mat, const Real* A, MatrixDim dst_d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda dst_index = i + j*dst_d.stride, src_index = i + j*src_stride;

if ( i < dst_d.cols && j < dst_d.rows ) {

Real a = mat[dst_index], b = A[src_index];

mat[dst_index] = (a > b ? a : b);

}

}

template<typename Real>

__global__

static void _vec_mul_elements(Real* v, const Real* a, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

if (i < dim)

v[i] = v[i] * a[i];

}

template<typename Real>

__global__

static void _mul_cols_vec(Real* mat, const Real* scale, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

mat[index] *= scale[i];

}

template<typename Real>

__global__

static void _mul_rows_vec(Real* mat, const Real* scale, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

mat[index] *= scale[j];

}

template<typename Real>

__global__

static void _mul_rows_group_mat(Real *y, const Real *x, MatrixDim d,

int src_stride, int group_size) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

int j = blockIdx.y * blockDim.y + threadIdx.y;

if (j < d.rows && i < d.cols ) {

int dst_index = i + j * d.stride;

int src_index = i / group_size + j * src_stride;

y[dst_index] *= x[src_index];

}

}

/// y is the derivative we will output; vec is the input we're computing

/// the group p-norm on, "norm" is the previously computed group p-norm.

template<typename Real>

__global__

static void _calc_pnorm_deriv(Real *deriv, const Real *vec, const Real *norm,

MatrixDim d, int src_stride, int group_size, Real power) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

int j = blockIdx.y * blockDim.y + threadIdx.y;

if (j < d.rows && i < d.cols ) {

int dst_index = i + j * d.stride,

src_index = i / group_size + j * src_stride;

Real vec_element = vec[dst_index], // this is the element of the original vector.

norm_element = norm[src_index]; // this is the pnorm

Real vec_element_sign = (vec_element > 0 ? 1 : -1);

Real ans;

if (norm_element <= 0.0) ans = 0.0; // The derivative is either zero or undefined at the origin.

else ans = vec_element_sign * pow(std::abs(vec_element), power - 1) *

pow(norm_element, 1 - power);

deriv[dst_index] = ans;

}

}

template<typename Real>

__global__

void _diff_group_pnorm(Real *id, const Real *iv, const Real *ov, const Real* od,

MatrixDim id_dim, int iv_stride, int ov_stride,

int od_stride, int group_size, Real power) {

const int j = blockIdx.x * blockDim.x + threadIdx.x;

if (j < id_dim.cols) {

const int grid_stride = gridDim.y * blockDim.y;

const int src_j = j / group_size;

int i = blockIdx.y * blockDim.y + threadIdx.y;

for (; i < id_dim.rows; i += grid_stride) {

const int iv_index = j + i * iv_stride;

Real iv_ij = iv[iv_index];

Real ans;

if (power == Real(2)) {

const int ov_index = src_j + i * ov_stride;

Real ov_ij = ov[ov_index];

ans = ov_ij <= 0.0 ? 0.0 : iv_ij / ov_ij;

} else if (power == Real(1)) {

Real iv_ij_sign = (iv_ij >= 0 ? 1 : -1);

ans = (iv_ij == Real(0) ? 0.0 : iv_ij_sign);

} else if (power

== (sizeof(Real) == sizeof(float) ? CUDART_INF_F : CUDART_INF)) {

const int ov_index = src_j + i * ov_stride;

Real ov_ij = ov[ov_index];

Real iv_ij_sign = (iv_ij >= 0 ? 1 : -1);

ans =

ov_ij <= 0.0 ?

0.0 : (iv_ij_sign * (abs(iv_ij) == ov_ij ? 1.0 : 0.0));

} else {

const int ov_index = src_j + i * ov_stride;

Real ov_ij = ov[ov_index];

Real iv_ij_sign = (iv_ij >= 0 ? 1 : -1);

if (ov_ij <= 0.0) {

ans = 0.0; // The derivative is either zero or undefined at the origin.

} else {

ans = iv_ij_sign * pow(std::abs(iv_ij), power - 1)

* pow(ov_ij, 1 - power);

}

}

const int od_index = src_j + i * od_stride;

const int id_index = j + i * id_dim.stride;

id[id_index] = ans * od[od_index];

}

}

}

/// deriv is the derivative we will output; vec is the input we're computing

/// the group max on, "maxv" is the previously computed group max.

template<typename Real>

__global__

static void _calc_group_max_deriv(Real *deriv, const Real *vec, const Real *maxv,

MatrixDim d, int src_stride, int group_size) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

int j = blockIdx.y * blockDim.y + threadIdx.y;

if (j < d.rows && i < d.cols ) {

int dst_index = i + j * d.stride,

src_index = i / group_size + j * src_stride;

Real vec_element = vec[dst_index], // this is the element of the original vector.

max_element = maxv[src_index]; // this is the max value

Real ans = (max_element == vec_element ? 1.0 : 0.0);

deriv[dst_index] = ans;

}

}

/// Set each element to y = (x == orig ? changed : x).

template<typename Real>

__global__

static void _replace_value(Real *vec, int dim, Real orig, Real changed) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

if (i < dim)

if (vec[i] == orig) vec[i] = changed;

}

template<typename Real>

__global__

static void _div_rows_vec(Real* mat, const Real* vec_div, MatrixDim d) {

const int32_cuda i = blockIdx.y * blockDim.y + threadIdx.y;

if (i < d.rows) {

const int32_cuda start = i * d.stride;

const Real scale = Real(1) / vec_div[i];

const int32_cuda grid_stride = blockDim.x * gridDim.x;

for (int32_cuda j = blockIdx.x * blockDim.x + threadIdx.x; j < d.cols; j +=

grid_stride) {

mat[start + j] *= scale;

}

}

}

template<typename Real>

__global__

static void _add_mat(Real alpha, const Real* src, Real* dst, MatrixDim d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x; // column index

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y; // row index

int32_cuda index = i + j * d.stride;

int32_cuda index_src = i + j * src_stride;

if (i < d.cols && j < d.rows)

dst[index] = alpha * src[index_src] + dst[index];

}

template<typename Real>

__global__

static void _add_mat_trans(Real alpha, const Real* src, Real* dst, MatrixDim d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j *d.stride;

int32_cuda index_src = j + i*src_stride;

if (i < d.cols && j < d.rows)

dst[index] = alpha*src[index_src] + dst[index];

}

template<typename Real>

__global__

static void _add_mat_blocks(Real alpha, const Real* src, int32_cuda num_row_blocks, int32_cuda num_col_blocks, Real* dst, MatrixDim d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j * d.stride;

int32_cuda index_src = i + j * src_stride;

if (i < d.cols && j < d.rows)

for (int32_cuda p = 0; p < num_row_blocks; p++) {

for (int32_cuda q = 0; q < num_col_blocks; q++) {

dst[index] = alpha * src[index_src + p * src_stride * d.rows + q * d.cols] + dst[index];

}

}

}

template<typename Real>

__global__

static void _add_mat_blocks_trans(Real alpha, const Real* src, int32_cuda num_row_blocks, int32_cuda num_col_blocks, Real* dst, MatrixDim d, int src_stride) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j * d.stride;

int32_cuda index_src = j + i * src_stride;

if (i < d.cols && j < d.rows)

for (int32_cuda p = 0; p < num_row_blocks; p++) {

for (int32_cuda q = 0; q < num_col_blocks; q++) {

dst[index] = alpha * src[index_src + p * src_stride * d.cols + q * d.rows] + dst[index];

}

}

}

template<typename Real>

__global__

static void _add_mat_mat_div_mat(const Real* A, const Real* B, const Real* C, Real* dst, MatrixDim d, int stride_a,

int stride_b, int stride_c) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride,

a_index = i + j*stride_a,

b_index = i + j*stride_b,

c_index = i + j*stride_c;

if (i < d.cols && j < d.rows)

if (C[c_index] == 0)

dst[index] = A[a_index];

else

dst[index] = A[a_index] * B[b_index] / C[c_index];

}

// Given a matrix input S (not packed!) and a lower-triangular matrix L,

// this function does S = beta S + alpha * L^T L. This is used in PSD matrix inversion.

// The i index is the row of the destination S and the j the column (although of

// course the output is symmetric so it doesn't matter in a sense). The main point

// of this is to make use of various symmetries and zero-ness.

template<typename Real>

__global__

static void _sy_add_tr2(Real alpha, Real beta, const Real *T, MatrixDim tdim, Real *S,

MatrixDim sdim) {

int i = blockIdx.x * blockDim.x + threadIdx.x;

int j = blockIdx.y * blockDim.y + threadIdx.y;

if (i >= sdim.rows || j > i) return;

// this thread computes the dot-product of the i'th column of

// L with the j'th column of L. The values we're multiplying

// are only nonzero for row-index k greater or equal to

// max(i, j), which equals i.

Real sum = 0.0;

for (int k = i; k < sdim.rows; k++) {

int i_index = i + tdim.stride * k,

j_index = j + tdim.stride * k;

sum += T[i_index] * T[j_index];

}

int output_index1 = i * sdim.stride + j,

output_index2 = j * sdim.stride + i;

S[output_index1] = alpha * sum + beta * S[output_index1];

S[output_index2] = alpha * sum + beta * S[output_index2];

}

template<typename Real>

__global__

static void _add_vec_to_cols(Real alpha, const Real* col, Real beta, Real* dst, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

dst[index] = alpha*col[j] + beta*dst[index];

}

template<typename Real>

__global__

static void _add_vec_to_rows(Real alpha, const Real* row, Real beta, Real* dst, MatrixDim d) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*d.stride;

if (i < d.cols && j < d.rows)

dst[index] = alpha*row[i] + beta*dst[index];

}

template<typename Real>

__global__

static void _apply_mask(Real* mat, const char* mask, MatrixDim dmat, MatrixDim dmask) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda index = i + j*dmat.stride;

int32_cuda index2 = i + j*dmask.stride;

if ( i < dmat.cols && j < dmat.rows )

if(mask[index2] == 0) mat[index] = 0;

}

template<typename Real>

__global__

static void _add_mat_diag_vec(Real alpha, Real *mat, MatrixDim mat_dim,

const Real *mat2, int mat2_row_stride, int mat2_col_stride,

const Real *vec, Real beta) {

int i = blockIdx.x * blockDim.x + threadIdx.x; // column index

int j = blockIdx.y * blockDim.y + threadIdx.y; // row index

int index = i + j * mat_dim.stride,

index2 = i * mat2_col_stride + j * mat2_row_stride;

if (j < mat_dim.rows && i < mat_dim.cols)

mat[index] = alpha * mat2[index2] * vec[i] + beta * mat[index];

}

template<typename Real>

__global__

static void _add_mat_mat_elements(Real *data, const Real *srcA_data, const Real *srcB_data, MatrixDim dim, int srcA_stride, int srcB_stride, Real alpha, Real beta) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda j = blockIdx.y * blockDim.y + threadIdx.y;

int32_cuda tgt_index = i + j*dim.stride;

int32_cuda srcA_index = i + j*srcA_stride;

int32_cuda srcB_index = i + j*srcB_stride;

if (i < dim.cols && j < dim.rows) {

data[tgt_index] = alpha * srcA_data[srcA_index] * srcB_data[srcB_index] + beta * data[tgt_index] ;

}

}

/*

* CuVector

*/

// very limited application!

template<typename Real>

__global__

static void _set_bias_params(Real* v, const Real* a, Real param_1, Real param_2, Real param_3, int* flag, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

if ( i < dim ) {

Real ratio = a[i] / param_3;

if ( ( ratio < 0.0 ) || ( ratio >= 1.01 )) {

*flag = 1;

return;

}

if ( ratio < param_1 ) {

Real factor = ((param_1/ratio) > param_2) ? param_2 : (param_1/ratio);

v[i] = v[i] / factor;

} else if ( ratio > param_1 ) {

Real factor = ((ratio/param_1) > param_2) ? param_2 : (ratio/param_1);

v[i] = v[i] * factor;

}

}

}

template<typename Real>

__global__

static void _copy_from_vec_df(double* v_out, const Real* v_in, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

// if (blockIdx.y > 0) return;

if (i < dim) {

v_out[i] = (double) v_in[i];

}

}

// This kernel writes a copy of the vector "v_in" to each row of the matrix

// "m_out". the dimension of v_in should be equal to the #columns of m_out.

template<typename Real>

__global__

static void _copy_rows_from_vec(Real* m_out, MatrixDim d, const Real* v_in) {

int i = blockIdx.x * blockDim.x + threadIdx.x; // column index.

int j = blockIdx.y * blockDim.y + threadIdx.y; // row index.

if (i < d.cols && j < d.rows) {

int index = i + j * d.stride;

m_out[index] = v_in[i];

}

}

template<typename Real>

__global__

static void _copy_from_vec_fd(float* v_out, const Real* v_in, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

// if (blockIdx.y > 0) return;

if (i < dim) {

v_out[i] = (float) v_in[i];

}

}

// _trace_mat_mat reduce the partial sum to value[blockIdx.y * gridDim.x + blockIdx.x]

// It use shared mem to transpose matrix B to ensure coalesced memory access

template<int TileDim, typename Real>

__global__

static void _trace_mat_mat(const Real* A, const Real* B, MatrixDim dA,

int B_stride, Real* value) {

// Reuse shared mem and make indexing easier. "+1" to avoid bank conflict

__shared__ union {

Real trans[TileDim][TileDim + 1];

Real sum[CU1DBLOCK];

} smem;

const int32_cuda tid = threadIdx.y * blockDim.x + threadIdx.x; // linear thread id;

const int32_cuda grid_height = gridDim.y * TileDim;

const int32_cuda ja = blockIdx.x * TileDim + threadIdx.x;

const int32_cuda ib = blockIdx.x * TileDim + threadIdx.y;

int32_cuda ia = blockIdx.y * TileDim + threadIdx.y;

int32_cuda jb = blockIdx.y * TileDim + threadIdx.x;

// Grid reduce

Real tsum = Real(0);

for (int32_cuda i0 = 0; i0 < dA.rows; i0 += grid_height) {

// Load from B, transpose the block and store in shared mem

if (jb < dA.rows) {

# pragma unroll

for (int i = 0; i < TileDim; i += CU1DBLOCK / TileDim) {

if (ib + i < dA.cols) {

smem.trans[threadIdx.x][threadIdx.y + i] =

B[(ib + i) * B_stride + jb];

}

}

}

__syncthreads();

// Load from A, sum up the product.

if (ja < dA.cols) {

# pragma unroll

for (int i = 0; i < TileDim; i += CU1DBLOCK / TileDim) {

if (ia + i < dA.rows) {

tsum += A[(ia + i) * dA.stride + ja]

* smem.trans[threadIdx.y + i][threadIdx.x];

}

}

}

__syncthreads();

ia += grid_height;

jb += grid_height;

}

smem.sum[tid] = tsum;

__syncthreads();

// Block reduce

# pragma unroll

for (int shift = CU1DBLOCK / 2; shift > warpSize; shift >>= 1) {

if (tid < shift)

smem.sum[tid] += smem.sum[tid + shift];

__syncthreads();

}

// Warp reduce. Implicitly synchronized within a warp.

if (tid < warpSize) {

# pragma unroll

for (int shift = warpSize; shift > 0; shift >>= 1) {

smem.sum[tid] += smem.sum[tid + shift];

}

}

// output 1 sum per thread block

if (tid == 0) {

value[blockIdx.y * gridDim.x + blockIdx.x] = smem.sum[0];

}

}

// _trace_mat_mat_trans reduce the partial sum to value[blockIdx.y * gridDim.x + blockIdx.x]

template<typename Real>

__global__

static void _trace_mat_mat_trans(const Real* A, const Real* B, MatrixDim dA, int B_stride, Real* value) {

__shared__ Real ssum[CU1DBLOCK];

const int32_cuda tid = threadIdx.y * blockDim.x + threadIdx.x; // linear thread id;

const int32_cuda j = blockIdx.x * blockDim.x + threadIdx.x;

const int32_cuda grid_height = gridDim.y * blockDim.y;

int32_cuda i = blockIdx.y * blockDim.y + threadIdx.y;

// Grid reduce

Real tsum = Real(0);

if (j < dA.cols) {

while (i < dA.rows) {

tsum += A[i * dA.stride + j] * B[i * B_stride + j];

i += grid_height;

}

}

ssum[tid] = tsum;

__syncthreads();

// Block reduce

# pragma unroll

for (int shift = CU1DBLOCK / 2; shift > warpSize; shift >>= 1) {

if (tid < shift)

ssum[tid] += ssum[tid + shift];

__syncthreads();

}

// Warp reduce. Implicitly synchronized within a warp.

if (tid < warpSize) {

# pragma unroll

for (int shift = warpSize; shift > 0; shift >>= 1) {

ssum[tid] += ssum[tid + shift];

}

}

// output 1 sum per thread block

if (tid == 0) {

value[blockIdx.y * gridDim.x + blockIdx.x] = ssum[0];

}

}

// Adds diag(M N) to v, where M and N are matrices. We supply row_stride and

// col_stride arguments for M and N, and swapping them allows us to transpose

// those matrices. Note: we imagine row-major indexing here, just like Kaldi

// and CBLAS (but unlike CUBLAS).

// This kernel expects the blockDim to be (CU1DBLOCK, 1) and the

// gridDim times CU1DBLOCK to be at least num-rows-of-v * threads_per_element.

// threads_per_element should be a power of 2.

template<typename Real>

__global__

static void _add_diag_mat_mat(

Real alpha, Real* v, int v_dim, const Real* M, int M_cols, int M_row_stride,

int M_col_stride, const Real *N, int N_row_stride, int N_col_stride,

int threads_per_element, Real beta) {

// we actually assume blockDim.x == CU1DBLOCK here.

// Each diagonal element of v is processed by "threads_per_element" threads.

__shared__ Real temp_data[CU1DBLOCK];

int i = blockIdx.x * blockDim.x + threadIdx.x;

int v_idx = i / threads_per_element, // v_idx is the index into v that we are supposed to

sub_idx = i % threads_per_element; // add to; 0 <= sub_idx < threads_per_element tells

// us which block of elements we sum up.

if (v_idx < v_dim) {

Real sum = 0.0;

for (int j = sub_idx; j < M_cols; j += threads_per_element) {

int M_index = v_idx * M_row_stride + j * M_col_stride,

N_index = j * N_row_stride + v_idx * N_col_stride;

sum += M[M_index] * N[N_index];

}

temp_data[threadIdx.x] = sum;

}

// start_idx = threadIdx.x - sub_idx; // start of the position in temp_data

// that we want to sum up.

// The following is a tree-based reduction of the elements of temp_data from

// start_idx to start_idx + threads_per_element - 1; our own index is "sub_idx".

__syncthreads();

int num_total_threads = threads_per_element;

while (num_total_threads > 1) {

int half_point = ((1 + num_total_threads) >> 1);

if (sub_idx < half_point) {

Real temp = 0.0;

if (sub_idx + half_point < num_total_threads) {

temp = temp_data[threadIdx.x + half_point];

}

temp_data[threadIdx.x] += temp;

}

__syncthreads();

num_total_threads = half_point;

}

if (sub_idx == 0 && v_idx < v_dim) {

v[v_idx] = beta * v[v_idx] + alpha * temp_data[threadIdx.x];

}

}

template<typename Real>

__global__

static void _add_vec_vec(Real alpha, Real* v, const Real* x, const Real* y, Real beta, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

// if (blockIdx.y > 0) return;

if (i < dim)

v[i] = alpha * x[i] * y[i] + beta * v[i];

}

template<typename Real>

__global__

static void _copy_col_from_mat_df(double* v, int col, const Real* mat, MatrixDim dmat, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda index = col + i * dmat.stride;

// if (blockIdx.y > 0) return;

if (i < dim)

v[i] = (double) mat[index];

}

template<typename Real>

__global__

static void _copy_col_from_mat_fd(float* v, int col, const Real* mat, MatrixDim dmat, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

int32_cuda index = col + i * dmat.stride;

// if (blockIdx.y > 0) return;

if (i < dim)

v[i] = (float) mat[index];

}

template<typename Real>

__global__

static void _vec_apply_exp(Real* v, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

// if (blockIdx.y > 0) return;

if (i < dim) {

v[i] = exp(v[i]);

}

}

template<typename Real>

__global__

static void _vec_apply_log(Real* v, Real* flag, int dim) {

int32_cuda i = blockIdx.x * blockDim.x + threadIdx.x;

// if (blockIdx.y > 0) return;

if (i < dim) {

if (v[i] < 0) {

*flag = 1;

return;

}

v[i] = log(v[i]);

}

}

template<typename Real>

__global__

static void _cuda_comp_obj_deriv(MatrixElement<Real> *x, int s, const Real* z, MatrixDim d, Real* z2, MatrixDim d2, Real* t) {

int i = threadIdx.x;

__shared__ Real tot_objf[CU1DBLOCK];