1 // Copyright (c) 2017-2022, 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 #include <libxsmm.h> 11 12 #include "ceed-xsmm.h" 13 14 //------------------------------------------------------------------------------ 15 // Tensor Contract Apply 16 //------------------------------------------------------------------------------ 17 static int CeedTensorContractApply_Xsmm(CeedTensorContract contract, CeedInt A, CeedInt B, CeedInt C, CeedInt J, const CeedScalar *restrict t, 18 CeedTransposeMode t_mode, const CeedInt add, const CeedScalar *restrict u, CeedScalar *restrict v) { 19 Ceed ceed; 20 21 CeedCallBackend(CeedTensorContractGetCeed(contract, &ceed)); 22 23 if (C == 1) { 24 // Build or query the required kernel 25 const int flags_t = LIBXSMM_GEMM_FLAGS(!t_mode ? 'T' : 'N', 'N'); 26 const int flags_ab = (!add) ? LIBXSMM_GEMM_FLAG_BETA_0 : LIBXSMM_BASIC_GEMM_FLAG_NONE; 27 const int flags = (flags_t | flags_ab); 28 const libxsmm_gemm_shape gemm_shape = (CEED_SCALAR_TYPE == CEED_SCALAR_FP64) 29 ? libxsmm_create_gemm_shape(J, A, B, !t_mode ? B : J, B, J, LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64, 30 LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64) 31 : libxsmm_create_gemm_shape(J, A, B, !t_mode ? B : J, B, J, LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32, 32 LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32); 33 const libxsmm_gemmfunction kernel = libxsmm_dispatch_gemm_v2(gemm_shape, (libxsmm_bitfield)(flags), (libxsmm_bitfield)LIBXSMM_GEMM_PREFETCH_NONE); 34 libxsmm_gemm_param gemm_param; 35 36 CeedCheck(kernel, ceed, CEED_ERROR_BACKEND, "LIBXSMM kernel failed to build."); 37 38 // Run kernel 39 gemm_param.a.primary = (CeedScalar *)&t[0]; 40 gemm_param.b.primary = (CeedScalar *)&u[0]; 41 gemm_param.c.primary = (CeedScalar *)&v[0]; 42 kernel(&gemm_param); 43 } else { 44 // Build or query the required kernel 45 const int flags_t = LIBXSMM_GEMM_FLAGS('N', t_mode ? 'T' : 'N'); 46 const int flags_ab = (!add) ? LIBXSMM_GEMM_FLAG_BETA_0 : LIBXSMM_BASIC_GEMM_FLAG_NONE; 47 const int flags = (flags_t | flags_ab); 48 const libxsmm_gemm_shape gemm_shape = (CEED_SCALAR_TYPE == CEED_SCALAR_FP64) 49 ? libxsmm_create_gemm_shape(C, J, B, C, !t_mode ? B : J, C, LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64, 50 LIBXSMM_DATATYPE_F64, LIBXSMM_DATATYPE_F64) 51 : libxsmm_create_gemm_shape(C, J, B, C, !t_mode ? B : J, C, LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32, 52 LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32); 53 const libxsmm_gemmfunction kernel = libxsmm_dispatch_gemm_v2(gemm_shape, (libxsmm_bitfield)(flags), (libxsmm_bitfield)LIBXSMM_GEMM_PREFETCH_NONE); 54 libxsmm_gemm_param gemm_param; 55 56 CeedCheck(kernel, ceed, CEED_ERROR_BACKEND, "LIBXSMM kernel failed to build."); 57 58 // Run kernel 59 gemm_param.b.primary = (CeedScalar *)&t[0]; 60 for (CeedInt a = 0; a < A; a++) { 61 gemm_param.a.primary = (CeedScalar *)&u[a * B * C]; 62 gemm_param.c.primary = (CeedScalar *)&v[a * J * C]; 63 kernel(&gemm_param); 64 } 65 } 66 return CEED_ERROR_SUCCESS; 67 } 68 69 //------------------------------------------------------------------------------ 70 // Tensor Contract Create 71 //------------------------------------------------------------------------------ 72 int CeedTensorContractCreate_Xsmm(CeedTensorContract contract) { 73 Ceed ceed; 74 75 CeedCallBackend(CeedTensorContractGetCeed(contract, &ceed)); 76 CeedCallBackend(CeedSetBackendFunction(ceed, "TensorContract", contract, "Apply", CeedTensorContractApply_Xsmm)); 77 return CEED_ERROR_SUCCESS; 78 } 79 80 //------------------------------------------------------------------------------ 81