   // For Upper/NoTrans and Lower/Trans process diagonal last
   k = min(n - bj, nb);
   while (k > 0) {
-    __syncthreads();
-
-#pragma unroll
-    for (int j = 0; j < nb; j += by)
-      a[threadIdx.x][j + threadIdx.y] = A[j * lda];
-
-    __syncthreads();
-
-    if (k < kb) break;
-
-// #pragma unroll
-//     for (int ll = 0; ll < kb; ll++) {
-//       saxpy(nb - ll, B[0], &a[ll][ll], &x[ll]);
-      saxpy(16, B[0], &a[ 0][ 0], &x[ 0]); B += ldb;
-      saxpy(15, B[0], &a[ 1][ 1], &x[ 1]); B += ldb;
-      saxpy(14, B[0], &a[ 2][ 2], &x[ 2]); B += ldb;
-      saxpy(13, B[0], &a[ 3][ 3], &x[ 3]); B += ldb;
-      saxpy(12, B[0], &a[ 4][ 4], &x[ 4]); B += ldb;
-      saxpy(11, B[0], &a[ 5][ 5], &x[ 5]); B += ldb;
-      saxpy(10, B[0], &a[ 6][ 6], &x[ 6]); B += ldb;
-      saxpy( 9, B[0], &a[ 7][ 7], &x[ 7]); B += ldb;
-      saxpy( 8, B[0], &a[ 8][ 8], &x[ 8]); B += ldb;
-      saxpy( 7, B[0], &a[ 9][ 9], &x[ 9]); B += ldb;
-      saxpy( 6, B[0], &a[10][10], &x[10]); B += ldb;
-      saxpy( 5, B[0], &a[11][11], &x[11]); B += ldb;
-      saxpy( 4, B[0], &a[12][12], &x[12]); B += ldb;
-      saxpy( 3, B[0], &a[13][13], &x[13]); B += ldb;
-      saxpy( 2, B[0], &a[14][14], &x[14]); B += ldb;
-      saxpy( 1, B[0], &a[15][15], &x[15]); B += ldb;
-//       B += ldb;
-//     }
-
-    __syncthreads();
+
+    __syncthreads();
+
+#pragma unroll
+    for (int j = 0; j < nb; j += by)
+      a[threadIdx.x][j + threadIdx.y] =
+        (diag != CBlasNonUnit && threadIdx.x == j + threadIdx.y) ? 1.0f : A[j * lda];
+
+    __syncthreads();
+
+    if (k < kb) break;
+
+    saxpy(16, B[0], &a[ 0][ 0], &x[ 0]); B += ldb;
+    saxpy(15, B[0], &a[ 1][ 1], &x[ 1]); B += ldb;
+    saxpy(14, B[0], &a[ 2][ 2], &x[ 2]); B += ldb;
+    saxpy(13, B[0], &a[ 3][ 3], &x[ 3]); B += ldb;
+    saxpy(12, B[0], &a[ 4][ 4], &x[ 4]); B += ldb;
+    saxpy(11, B[0], &a[ 5][ 5], &x[ 5]); B += ldb;
+    saxpy(10, B[0], &a[ 6][ 6], &x[ 6]); B += ldb;
+    saxpy( 9, B[0], &a[ 7][ 7], &x[ 7]); B += ldb;
+    saxpy( 8, B[0], &a[ 8][ 8], &x[ 8]); B += ldb;
+    saxpy( 7, B[0], &a[ 9][ 9], &x[ 9]); B += ldb;
+    saxpy( 6, B[0], &a[10][10], &x[10]); B += ldb;
+    saxpy( 5, B[0], &a[11][11], &x[11]); B += ldb;
+    saxpy( 4, B[0], &a[12][12], &x[12]); B += ldb;
+    saxpy( 3, B[0], &a[13][13], &x[13]); B += ldb;
+    saxpy( 2, B[0], &a[14][14], &x[14]); B += ldb;
+    saxpy( 1, B[0], &a[15][15], &x[15]); B += ldb;
 
     A += kb;
     k -= kb;
   }
 
-//     for (int ll = 0; ll < k; ll++) {
-//       saxpy(nb - ll, B[0], &a[ll][ll], &x[ll]);
-//       B += ldb;
-//     }
   if (k > 0) { saxpy(16, B[0], &a[ 0][ 0], &x[ 0]); B += ldb;
   if (k > 1) { saxpy(15, B[0], &a[ 1][ 1], &x[ 1]); B += ldb;
   if (k > 2) { saxpy(14, B[0], &a[ 2][ 2], &x[ 2]); B += ldb;


template void strmmRUN<CBlasUnit,    64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
template void strmmRUN<CBlasNonUnit, 64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
// template void strmm2R<CBlasUpper, CBlasTrans,   CBlasUnit,    64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
// template void strmm2R<CBlasUpper, CBlasTrans,   CBlasNonUnit, 64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
// template void strmm2R<CBlasLower, CBlasNoTrans, CBlasUnit,    64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
// template void strmm2R<CBlasLower, CBlasNoTrans, CBlasNonUnit, 64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
// template void strmm2R<CBlasLower, CBlasTrans,   CBlasUnit,    64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);
// template void strmm2R<CBlasLower, CBlasTrans,   CBlasNonUnit, 64, 16, 16, 16,  4>(int, int, float, const float * __restrict__, int, const float * __restrict__, int, float * __restrict__, int);

          unsigned int mb, unsigned int nb, unsigned int kb,
          unsigned int bx, unsigned int by>
__global__ void strmmRUN(int m, int n,
                         float alpha, const float * __restrict__ A, int lda,
                         const float * __restrict__ B, int ldb,
                         float * __restrict__ X, int ldx) {

  const int bi = blockIdx.x * mb;       // Starting row of block of X
  const int bj = blockIdx.y * nb;       // Starting column of block of X
  int ti = threadIdx.y * bx + threadIdx.x;

  A += (bj + threadIdx.y) * lda + threadIdx.x;
  B += bi + ti;
