jit-source/magma/magma-common-tensor.h

1 // Copyright (c) 2017-2026, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
4 // SPDX-License-Identifier: BSD-2-Clause
12 #include "magma-common-defs.h"
15 // read U or V of a 1D element into shared memory sU[][] or sV[][] --  for all components
28 // write V of a 1D element into global memory from sV[][] --  for all components
40 // sum into V of a 1D element into global memory from sV[][] --  for all components
52 // read U of a 2D element into registers rU[][][] --  for all components of a single dim
53 // dU is assumed to be offset by elem-stride and dim-stride
61   // vec 0  : [u0, u1, u2, ... u_(P-1)] -- contiguous in memory  in read_U_2d()
62   // vec 1  : [u0, u1, u2, ... u_(P-1)] -- contiguous in memory  in read_U_2d()
64   // vec P-1: [u0, u1, u2, ... u_(P-1)] -- contiguous in memory  in read_U_2d()
67   // thread 0 to hold all of vec0 in registers, and  in read_U_2d()
68   // thread 1 to hold all of vec1 in registers, and and so on  in read_U_2d()
89 // read V of a 2D element into registers rV[][][] --  for all components of a single dim
90 // dV is assumed to be offset by elem-stride and dim-stride
106 // write V of a 2D element from registers rV[][][] to global memory --  for all components of a sin…
107 // dV is assumed to be offset by elem-stride and dim-stride
123 // sum into V of a 2D element from registers rV[][][] to global memory --  for all components of a …
124 // dV is assumed to be offset by elem-stride and dim-stride
140 // read U of a 3D element into registers rU[][][] --  for all components of a single dim
141 // dU is assumed to be offset by elem-stride and dim-stride
149   // vec 0    : [u0, u1, u2, ... u_(P-1)] -- contiguous in memory  in read_U_3d()
150   // vec 1    : [u0, u1, u2, ... u_(P-1)] -- contiguous in memory  in read_U_3d()
152   // vec P^2-1: [u0, u1, u2, ... u_(P-1)] -- contiguous in memory  in read_U_3d()
155   // thread 0 to hold all of vec0 in registers, and  in read_U_3d()
156   // thread 1 to hold all of vec1 in registers, and and so on  in read_U_3d()
177 // read V of a 3D element into registers rV[][][] --  for all components of a single dim
178 // dV is assumed to be offset by elem-stride and dim-stride
194 // write V of a 3D element from registers rV[][][] to global memory --  for all components of a sin…
195 // dV is assumed to point directly to the element (i.e. already offset by elem-stride)
211 // sum into V of a 3D element from registers rV[][][] to global memory --  for all components of a …
212 // dV is assumed to point directly to the element (i.e. already offset by elem-stride)
228 // reads T (no-trans) into shared memory