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.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED: http://github.com/ceed
7
8 #include <ceed.h>
9 #include <ceed/backend.h>
10
11 #include "ceed-opt.h"
12
13 //------------------------------------------------------------------------------
14 // Tensor Contract Core loop
15 //------------------------------------------------------------------------------
CeedTensorContractApply_Core_Opt(CeedTensorContract contract,CeedInt A,CeedInt B,CeedInt C,CeedInt J,const CeedScalar * restrict t,CeedTransposeMode t_mode,const CeedInt add,const CeedScalar * restrict u,CeedScalar * restrict v)16 static inline int CeedTensorContractApply_Core_Opt(CeedTensorContract contract, CeedInt A, CeedInt B, CeedInt C, CeedInt J,
17 const CeedScalar *restrict t, CeedTransposeMode t_mode, const CeedInt add,
18 const CeedScalar *restrict u, CeedScalar *restrict v) {
19 CeedInt t_stride_0 = B, t_stride_1 = 1;
20
21 if (t_mode == CEED_TRANSPOSE) {
22 t_stride_0 = 1;
23 t_stride_1 = J;
24 }
25
26 for (CeedInt a = 0; a < A; a++) {
27 for (CeedInt b = 0; b < B; b++) {
28 for (CeedInt j = 0; j < J; j++) {
29 CeedScalar tq = t[j * t_stride_0 + b * t_stride_1];
30 for (CeedInt c = 0; c < C; c++) v[(a * J + j) * C + c] += tq * u[(a * B + b) * C + c];
31 }
32 }
33 }
34 return CEED_ERROR_SUCCESS;
35 }
36
37 //------------------------------------------------------------------------------
38 // Tensor Contract Apply
39 //------------------------------------------------------------------------------
CeedTensorContractApply_Opt(CeedTensorContract contract,CeedInt A,CeedInt B,CeedInt C,CeedInt J,const CeedScalar * restrict t,CeedTransposeMode t_mode,const CeedInt add,const CeedScalar * restrict u,CeedScalar * restrict v)40 static int CeedTensorContractApply_Opt(CeedTensorContract contract, CeedInt A, CeedInt B, CeedInt C, CeedInt J, const CeedScalar *restrict t,
41 CeedTransposeMode t_mode, const CeedInt add, const CeedScalar *restrict u, CeedScalar *restrict v) {
42 if (!add) {
43 for (CeedInt q = 0; q < A * J * C; q++) v[q] = (CeedScalar)0.0;
44 }
45
46 if (C == 1) return CeedTensorContractApply_Core_Opt(contract, A, B, 1, J, t, t_mode, add, u, v);
47 else return CeedTensorContractApply_Core_Opt(contract, A, B, C, J, t, t_mode, add, u, v);
48 return CEED_ERROR_SUCCESS;
49 }
50
51 //------------------------------------------------------------------------------
52 // Tensor Contract Create
53 //------------------------------------------------------------------------------
CeedTensorContractCreate_Opt(CeedTensorContract contract)54 int CeedTensorContractCreate_Opt(CeedTensorContract contract) {
55 CeedCallBackend(CeedSetBackendFunction(CeedTensorContractReturnCeed(contract), "TensorContract", contract, "Apply", CeedTensorContractApply_Opt));
56 return CEED_ERROR_SUCCESS;
57 }
58
59 //------------------------------------------------------------------------------
60