1 // Copyright (c) 2017-2024, 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 <ceed/jit-source/hip/hip-types.h> 11 #include <stddef.h> 12 13 #include "../hip/ceed-hip-common.h" 14 #include "../hip/ceed-hip-compile.h" 15 #include "ceed-hip-gen-operator-build.h" 16 #include "ceed-hip-gen.h" 17 18 //------------------------------------------------------------------------------ 19 // Destroy operator 20 //------------------------------------------------------------------------------ 21 static int CeedOperatorDestroy_Hip_gen(CeedOperator op) { 22 CeedOperator_Hip_gen *impl; 23 24 CeedCallBackend(CeedOperatorGetData(op, &impl)); 25 CeedCallBackend(CeedFree(&impl)); 26 return CEED_ERROR_SUCCESS; 27 } 28 29 //------------------------------------------------------------------------------ 30 // Apply and add to output 31 //------------------------------------------------------------------------------ 32 static int CeedOperatorApplyAdd_Hip_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) { 33 Ceed ceed; 34 CeedInt num_elem, num_input_fields, num_output_fields; 35 CeedEvalMode eval_mode; 36 CeedVector output_vecs[CEED_FIELD_MAX] = {NULL}; 37 CeedQFunctionField *qf_input_fields, *qf_output_fields; 38 CeedQFunction_Hip_gen *qf_data; 39 CeedQFunction qf; 40 CeedOperatorField *op_input_fields, *op_output_fields; 41 CeedOperator_Hip_gen *data; 42 43 // Check for tensor-product bases 44 { 45 bool has_tensor_bases; 46 47 CeedCallBackend(CeedOperatorHasTensorBases(op, &has_tensor_bases)); 48 // -- Fallback to ref if not all bases are tensor-product 49 if (!has_tensor_bases) { 50 CeedOperator op_fallback; 51 52 CeedDebug256(CeedOperatorReturnCeed(op), CEED_DEBUG_COLOR_SUCCESS, "Falling back to /gpu/hip/ref CeedOperator due to non-tensor bases"); 53 CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback)); 54 CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request)); 55 return CEED_ERROR_SUCCESS; 56 } 57 } 58 59 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 60 CeedCallBackend(CeedOperatorGetData(op, &data)); 61 CeedCallBackend(CeedOperatorGetQFunction(op, &qf)); 62 CeedCallBackend(CeedQFunctionGetData(qf, &qf_data)); 63 CeedCallBackend(CeedOperatorGetNumElements(op, &num_elem)); 64 CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields)); 65 CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields)); 66 67 // Creation of the operator 68 CeedCallBackend(CeedOperatorBuildKernel_Hip_gen(op)); 69 70 // Input vectors 71 for (CeedInt i = 0; i < num_input_fields; i++) { 72 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 73 if (eval_mode == CEED_EVAL_WEIGHT) { // Skip 74 data->fields.inputs[i] = NULL; 75 } else { 76 CeedVector vec; 77 78 // Get input vector 79 CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec)); 80 if (vec == CEED_VECTOR_ACTIVE) vec = input_vec; 81 CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->fields.inputs[i])); 82 } 83 } 84 85 // Output vectors 86 for (CeedInt i = 0; i < num_output_fields; i++) { 87 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 88 if (eval_mode == CEED_EVAL_WEIGHT) { // Skip 89 data->fields.outputs[i] = NULL; 90 } else { 91 CeedVector vec; 92 93 // Get output vector 94 CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec)); 95 if (vec == CEED_VECTOR_ACTIVE) vec = output_vec; 96 output_vecs[i] = vec; 97 // Check for multiple output modes 98 CeedInt index = -1; 99 for (CeedInt j = 0; j < i; j++) { 100 if (vec == output_vecs[j]) { 101 index = j; 102 break; 103 } 104 } 105 if (index == -1) { 106 CeedCallBackend(CeedVectorGetArray(vec, CEED_MEM_DEVICE, &data->fields.outputs[i])); 107 } else { 108 data->fields.outputs[i] = data->fields.outputs[index]; 109 } 110 } 111 } 112 113 // Get context data 114 CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &qf_data->d_c)); 115 116 // Apply operator 117 void *opargs[] = {(void *)&num_elem, &qf_data->d_c, &data->indices, &data->fields, &data->B, &data->G, &data->W}; 118 const CeedInt dim = data->dim; 119 const CeedInt Q_1d = data->Q_1d; 120 const CeedInt P_1d = data->max_P_1d; 121 const CeedInt thread_1d = CeedIntMax(Q_1d, P_1d); 122 CeedInt block_sizes[3]; 123 124 CeedCallBackend(BlockGridCalculate_Hip_gen(dim, num_elem, P_1d, Q_1d, block_sizes)); 125 if (dim == 1) { 126 CeedInt grid = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0); 127 CeedInt sharedMem = block_sizes[2] * thread_1d * sizeof(CeedScalar); 128 129 CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs)); 130 } else if (dim == 2) { 131 CeedInt grid = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0); 132 CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar); 133 134 CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs)); 135 } else if (dim == 3) { 136 CeedInt grid = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0); 137 CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar); 138 139 CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs)); 140 } 141 142 // Restore input arrays 143 for (CeedInt i = 0; i < num_input_fields; i++) { 144 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 145 if (eval_mode == CEED_EVAL_WEIGHT) { // Skip 146 } else { 147 CeedVector vec; 148 149 CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec)); 150 if (vec == CEED_VECTOR_ACTIVE) vec = input_vec; 151 CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->fields.inputs[i])); 152 } 153 } 154 155 // Restore output arrays 156 for (CeedInt i = 0; i < num_output_fields; i++) { 157 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 158 if (eval_mode == CEED_EVAL_WEIGHT) { // Skip 159 } else { 160 CeedVector vec; 161 162 CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec)); 163 if (vec == CEED_VECTOR_ACTIVE) vec = output_vec; 164 // Check for multiple output modes 165 CeedInt index = -1; 166 for (CeedInt j = 0; j < i; j++) { 167 if (vec == output_vecs[j]) { 168 index = j; 169 break; 170 } 171 } 172 if (index == -1) { 173 CeedCallBackend(CeedVectorRestoreArray(vec, &data->fields.outputs[i])); 174 } 175 } 176 } 177 178 // Restore context data 179 CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &qf_data->d_c)); 180 CeedCallBackend(CeedDestroy(&ceed)); 181 CeedCallBackend(CeedQFunctionDestroy(&qf)); 182 return CEED_ERROR_SUCCESS; 183 } 184 185 //------------------------------------------------------------------------------ 186 // Create operator 187 //------------------------------------------------------------------------------ 188 int CeedOperatorCreate_Hip_gen(CeedOperator op) { 189 Ceed ceed; 190 CeedOperator_Hip_gen *impl; 191 192 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 193 CeedCallBackend(CeedCalloc(1, &impl)); 194 CeedCallBackend(CeedOperatorSetData(op, impl)); 195 CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Hip_gen)); 196 CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Hip_gen)); 197 CeedCallBackend(CeedDestroy(&ceed)); 198 return CEED_ERROR_SUCCESS; 199 } 200 201 //------------------------------------------------------------------------------ 202