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-tools.h> 11 #include <stdbool.h> 12 #include <stddef.h> 13 #include <string.h> 14 #include <hip/hip_runtime.h> 15 16 #include "../hip/ceed-hip-common.h" 17 #include "../hip/ceed-hip-compile.h" 18 #include "ceed-hip-ref.h" 19 20 //------------------------------------------------------------------------------ 21 // Compile restriction kernels 22 //------------------------------------------------------------------------------ 23 static inline int CeedElemRestrictionSetupCompile_Hip(CeedElemRestriction rstr) { 24 Ceed ceed; 25 bool is_deterministic; 26 char *restriction_kernel_source; 27 const char *restriction_kernel_path; 28 CeedInt num_elem, num_comp, elem_size, comp_stride; 29 CeedRestrictionType rstr_type; 30 CeedElemRestriction_Hip *impl; 31 32 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 33 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed)); 34 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type)); 35 CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem)); 36 CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp)); 37 CeedCallBackend(CeedElemRestrictionGetCompStride(rstr, &comp_stride)); 38 if (rstr_type == CEED_RESTRICTION_POINTS) { 39 CeedCallBackend(CeedElemRestrictionGetMaxPointsInElement(rstr, &elem_size)); 40 } else { 41 CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size)); 42 } 43 is_deterministic = impl->d_l_vec_indices != NULL; 44 45 // Compile HIP kernels 46 switch (rstr_type) { 47 case CEED_RESTRICTION_STRIDED: { 48 bool has_backend_strides; 49 CeedInt strides[3] = {1, num_elem * elem_size, elem_size}; 50 51 CeedCallBackend(CeedElemRestrictionHasBackendStrides(rstr, &has_backend_strides)); 52 if (!has_backend_strides) { 53 CeedCallBackend(CeedElemRestrictionGetStrides(rstr, strides)); 54 } 55 56 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-strided.h", &restriction_kernel_path)); 57 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n"); 58 CeedCallBackend(CeedLoadSourceToBuffer(ceed, restriction_kernel_path, &restriction_kernel_source)); 59 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n"); 60 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem, 61 "RSTR_NUM_COMP", num_comp, "RSTR_STRIDE_NODES", strides[0], "RSTR_STRIDE_COMP", strides[1], "RSTR_STRIDE_ELEM", 62 strides[2])); 63 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "StridedNoTranspose", &impl->ApplyNoTranspose)); 64 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "StridedTranspose", &impl->ApplyTranspose)); 65 } break; 66 case CEED_RESTRICTION_POINTS: 67 case CEED_RESTRICTION_STANDARD: { 68 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-offset.h", &restriction_kernel_path)); 69 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n"); 70 CeedCallBackend(CeedLoadSourceToBuffer(ceed, restriction_kernel_path, &restriction_kernel_source)); 71 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n"); 72 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem, 73 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride, 74 "USE_DETERMINISTIC", is_deterministic ? 1 : 0)); 75 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyNoTranspose)); 76 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyTranspose)); 77 } break; 78 case CEED_RESTRICTION_ORIENTED: { 79 const char *offset_kernel_path; 80 char **file_paths = NULL; 81 CeedInt num_file_paths = 0; 82 83 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-oriented.h", &restriction_kernel_path)); 84 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n"); 85 CeedCallBackend(CeedLoadSourceAndInitializeBuffer(ceed, restriction_kernel_path, &num_file_paths, &file_paths, &restriction_kernel_source)); 86 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-offset.h", &offset_kernel_path)); 87 CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, offset_kernel_path, &num_file_paths, &file_paths, &restriction_kernel_source)); 88 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n"); 89 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem, 90 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride, 91 "USE_DETERMINISTIC", is_deterministic ? 1 : 0)); 92 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OrientedNoTranspose", &impl->ApplyNoTranspose)); 93 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyUnsignedNoTranspose)); 94 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OrientedTranspose", &impl->ApplyTranspose)); 95 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyUnsignedTranspose)); 96 // Cleanup 97 CeedCallBackend(CeedFree(&offset_kernel_path)); 98 for (CeedInt i = 0; i < num_file_paths; i++) CeedCall(CeedFree(&file_paths[i])); 99 CeedCall(CeedFree(&file_paths)); 100 } break; 101 case CEED_RESTRICTION_CURL_ORIENTED: { 102 const char *offset_kernel_path; 103 char **file_paths = NULL; 104 CeedInt num_file_paths = 0; 105 106 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-curl-oriented.h", &restriction_kernel_path)); 107 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n"); 108 CeedCallBackend(CeedLoadSourceAndInitializeBuffer(ceed, restriction_kernel_path, &num_file_paths, &file_paths, &restriction_kernel_source)); 109 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-offset.h", &offset_kernel_path)); 110 CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, offset_kernel_path, &num_file_paths, &file_paths, &restriction_kernel_source)); 111 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n"); 112 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem, 113 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride, 114 "USE_DETERMINISTIC", is_deterministic ? 1 : 0)); 115 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedNoTranspose", &impl->ApplyNoTranspose)); 116 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedUnsignedNoTranspose", &impl->ApplyUnsignedNoTranspose)); 117 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyUnorientedNoTranspose)); 118 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedTranspose", &impl->ApplyTranspose)); 119 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedUnsignedTranspose", &impl->ApplyUnsignedTranspose)); 120 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyUnorientedTranspose)); 121 // Cleanup 122 CeedCallBackend(CeedFree(&offset_kernel_path)); 123 for (CeedInt i = 0; i < num_file_paths; i++) CeedCall(CeedFree(&file_paths[i])); 124 CeedCall(CeedFree(&file_paths)); 125 } break; 126 } 127 CeedCallBackend(CeedFree(&restriction_kernel_path)); 128 CeedCallBackend(CeedFree(&restriction_kernel_source)); 129 return CEED_ERROR_SUCCESS; 130 } 131 132 //------------------------------------------------------------------------------ 133 // Core apply restriction code 134 //------------------------------------------------------------------------------ 135 static inline int CeedElemRestrictionApply_Hip_Core(CeedElemRestriction rstr, CeedTransposeMode t_mode, bool use_signs, bool use_orients, 136 CeedVector u, CeedVector v, CeedRequest *request) { 137 Ceed ceed; 138 CeedRestrictionType rstr_type; 139 const CeedScalar *d_u; 140 CeedScalar *d_v; 141 CeedElemRestriction_Hip *impl; 142 143 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 144 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed)); 145 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type)); 146 147 // Assemble kernel if needed 148 if (!impl->module) { 149 CeedCallBackend(CeedElemRestrictionSetupCompile_Hip(rstr)); 150 } 151 152 // Get vectors 153 CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u)); 154 if (t_mode == CEED_TRANSPOSE) { 155 // Sum into for transpose mode, e-vec to l-vec 156 CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v)); 157 } else { 158 // Overwrite for notranspose mode, l-vec to e-vec 159 CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v)); 160 } 161 162 // Restrict 163 if (t_mode == CEED_NOTRANSPOSE) { 164 // L-vector -> E-vector 165 CeedInt elem_size; 166 167 CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size)); 168 const CeedInt block_size = elem_size < 256 ? (elem_size > 64 ? elem_size : 64) : 256; 169 const CeedInt grid = CeedDivUpInt(impl->num_nodes, block_size); 170 171 switch (rstr_type) { 172 case CEED_RESTRICTION_STRIDED: { 173 void *args[] = {&d_u, &d_v}; 174 175 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args)); 176 } break; 177 case CEED_RESTRICTION_POINTS: 178 case CEED_RESTRICTION_STANDARD: { 179 void *args[] = {&impl->d_offsets, &d_u, &d_v}; 180 181 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args)); 182 } break; 183 case CEED_RESTRICTION_ORIENTED: { 184 if (use_signs) { 185 void *args[] = {&impl->d_offsets, &impl->d_orients, &d_u, &d_v}; 186 187 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args)); 188 } else { 189 void *args[] = {&impl->d_offsets, &d_u, &d_v}; 190 191 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedNoTranspose, grid, block_size, args)); 192 } 193 } break; 194 case CEED_RESTRICTION_CURL_ORIENTED: { 195 if (use_signs && use_orients) { 196 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v}; 197 198 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args)); 199 } else if (use_orients) { 200 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v}; 201 202 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedNoTranspose, grid, block_size, args)); 203 } else { 204 void *args[] = {&impl->d_offsets, &d_u, &d_v}; 205 206 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedNoTranspose, grid, block_size, args)); 207 } 208 } break; 209 } 210 } else { 211 // E-vector -> L-vector 212 const bool is_deterministic = impl->d_l_vec_indices != NULL; 213 const CeedInt block_size = 64; 214 const CeedInt grid = CeedDivUpInt(impl->num_nodes, block_size); 215 216 switch (rstr_type) { 217 case CEED_RESTRICTION_STRIDED: { 218 void *args[] = {&d_u, &d_v}; 219 220 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 221 } break; 222 case CEED_RESTRICTION_POINTS: 223 case CEED_RESTRICTION_STANDARD: { 224 if (!is_deterministic) { 225 void *args[] = {&impl->d_offsets, &d_u, &d_v}; 226 227 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 228 } else { 229 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &d_u, &d_v}; 230 231 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 232 } 233 } break; 234 case CEED_RESTRICTION_ORIENTED: { 235 if (use_signs) { 236 if (!is_deterministic) { 237 void *args[] = {&impl->d_offsets, &impl->d_orients, &d_u, &d_v}; 238 239 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 240 } else { 241 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &impl->d_orients, &d_u, &d_v}; 242 243 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 244 } 245 } else { 246 if (!is_deterministic) { 247 void *args[] = {&impl->d_offsets, &d_u, &d_v}; 248 249 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args)); 250 } else { 251 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &d_u, &d_v}; 252 253 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args)); 254 } 255 } 256 } break; 257 case CEED_RESTRICTION_CURL_ORIENTED: { 258 if (use_signs && use_orients) { 259 if (!is_deterministic) { 260 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v}; 261 262 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 263 } else { 264 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &impl->d_curl_orients, &d_u, &d_v}; 265 266 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args)); 267 } 268 } else if (use_orients) { 269 if (!is_deterministic) { 270 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v}; 271 272 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args)); 273 } else { 274 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &impl->d_curl_orients, &d_u, &d_v}; 275 276 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args)); 277 } 278 } else { 279 if (!is_deterministic) { 280 void *args[] = {&impl->d_offsets, &d_u, &d_v}; 281 282 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedTranspose, grid, block_size, args)); 283 } else { 284 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &d_u, &d_v}; 285 286 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedTranspose, grid, block_size, args)); 287 } 288 } 289 } break; 290 } 291 } 292 293 if (request != CEED_REQUEST_IMMEDIATE && request != CEED_REQUEST_ORDERED) *request = NULL; 294 295 // Restore arrays 296 CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u)); 297 CeedCallBackend(CeedVectorRestoreArray(v, &d_v)); 298 return CEED_ERROR_SUCCESS; 299 } 300 301 //------------------------------------------------------------------------------ 302 // Apply restriction 303 //------------------------------------------------------------------------------ 304 static int CeedElemRestrictionApply_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v, CeedRequest *request) { 305 return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, true, true, u, v, request); 306 } 307 308 //------------------------------------------------------------------------------ 309 // Apply unsigned restriction 310 //------------------------------------------------------------------------------ 311 static int CeedElemRestrictionApplyUnsigned_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v, 312 CeedRequest *request) { 313 return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, false, true, u, v, request); 314 } 315 316 //------------------------------------------------------------------------------ 317 // Apply unoriented restriction 318 //------------------------------------------------------------------------------ 319 static int CeedElemRestrictionApplyUnoriented_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v, 320 CeedRequest *request) { 321 return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, false, false, u, v, request); 322 } 323 324 //------------------------------------------------------------------------------ 325 // Get offsets 326 //------------------------------------------------------------------------------ 327 static int CeedElemRestrictionGetOffsets_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const CeedInt **offsets) { 328 CeedElemRestriction_Hip *impl; 329 CeedRestrictionType rstr_type; 330 331 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 332 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type)); 333 switch (mem_type) { 334 case CEED_MEM_HOST: 335 *offsets = rstr_type == CEED_RESTRICTION_POINTS ? impl->h_offsets_at_points : impl->h_offsets; 336 break; 337 case CEED_MEM_DEVICE: 338 *offsets = rstr_type == CEED_RESTRICTION_POINTS ? impl->d_offsets_at_points : impl->d_offsets; 339 break; 340 } 341 return CEED_ERROR_SUCCESS; 342 } 343 344 //------------------------------------------------------------------------------ 345 // Get orientations 346 //------------------------------------------------------------------------------ 347 static int CeedElemRestrictionGetOrientations_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const bool **orients) { 348 CeedElemRestriction_Hip *impl; 349 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 350 351 switch (mem_type) { 352 case CEED_MEM_HOST: 353 *orients = impl->h_orients; 354 break; 355 case CEED_MEM_DEVICE: 356 *orients = impl->d_orients; 357 break; 358 } 359 return CEED_ERROR_SUCCESS; 360 } 361 362 //------------------------------------------------------------------------------ 363 // Get curl-conforming orientations 364 //------------------------------------------------------------------------------ 365 static int CeedElemRestrictionGetCurlOrientations_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const CeedInt8 **curl_orients) { 366 CeedElemRestriction_Hip *impl; 367 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 368 369 switch (mem_type) { 370 case CEED_MEM_HOST: 371 *curl_orients = impl->h_curl_orients; 372 break; 373 case CEED_MEM_DEVICE: 374 *curl_orients = impl->d_curl_orients; 375 break; 376 } 377 return CEED_ERROR_SUCCESS; 378 } 379 380 //------------------------------------------------------------------------------ 381 // Get offset for padded AtPoints E-layout 382 //------------------------------------------------------------------------------ 383 static int CeedElemRestrictionGetAtPointsElementOffset_Hip(CeedElemRestriction rstr, CeedInt elem, CeedSize *elem_offset) { 384 CeedInt layout[3]; 385 386 CeedCallBackend(CeedElemRestrictionGetELayout(rstr, layout)); 387 *elem_offset = 0 * layout[0] + 0 * layout[1] + elem * layout[2]; 388 return CEED_ERROR_SUCCESS; 389 } 390 391 //------------------------------------------------------------------------------ 392 // Destroy restriction 393 //------------------------------------------------------------------------------ 394 static int CeedElemRestrictionDestroy_Hip(CeedElemRestriction rstr) { 395 Ceed ceed; 396 CeedElemRestriction_Hip *impl; 397 398 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 399 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed)); 400 if (impl->module) { 401 CeedCallHip(ceed, hipModuleUnload(impl->module)); 402 } 403 CeedCallBackend(CeedFree(&impl->h_offsets_owned)); 404 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_offsets_owned)); 405 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_t_offsets)); 406 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_t_indices)); 407 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_l_vec_indices)); 408 CeedCallBackend(CeedFree(&impl->h_orients_owned)); 409 CeedCallHip(ceed, hipFree((bool *)impl->d_orients_owned)); 410 CeedCallBackend(CeedFree(&impl->h_curl_orients_owned)); 411 CeedCallHip(ceed, hipFree((CeedInt8 *)impl->d_curl_orients_owned)); 412 CeedCallBackend(CeedFree(&impl->h_offsets_at_points_owned)); 413 CeedCallHip(ceed, hipFree((CeedInt8 *)impl->d_offsets_at_points_owned)); 414 CeedCallBackend(CeedFree(&impl)); 415 return CEED_ERROR_SUCCESS; 416 } 417 418 //------------------------------------------------------------------------------ 419 // Create transpose offsets and indices 420 //------------------------------------------------------------------------------ 421 static int CeedElemRestrictionOffset_Hip(const CeedElemRestriction rstr, const CeedInt elem_size, const CeedInt *indices) { 422 Ceed ceed; 423 bool *is_node; 424 CeedSize l_size; 425 CeedInt num_elem, num_comp, num_nodes = 0; 426 CeedInt *ind_to_offset, *l_vec_indices, *t_offsets, *t_indices; 427 CeedRestrictionType rstr_type; 428 CeedElemRestriction_Hip *impl; 429 430 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed)); 431 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl)); 432 CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem)); 433 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type)); 434 CeedCallBackend(CeedElemRestrictionGetLVectorSize(rstr, &l_size)); 435 CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp)); 436 const CeedInt size_indices = num_elem * elem_size; 437 438 // Count num_nodes 439 CeedCallBackend(CeedCalloc(l_size, &is_node)); 440 441 for (CeedInt i = 0; i < size_indices; i++) is_node[indices[i]] = 1; 442 for (CeedInt i = 0; i < l_size; i++) num_nodes += is_node[i]; 443 impl->num_nodes = num_nodes; 444 445 // L-vector offsets array 446 CeedCallBackend(CeedCalloc(l_size, &ind_to_offset)); 447 CeedCallBackend(CeedCalloc(num_nodes, &l_vec_indices)); 448 for (CeedInt i = 0, j = 0; i < l_size; i++) { 449 if (is_node[i]) { 450 l_vec_indices[j] = i; 451 ind_to_offset[i] = j++; 452 } 453 } 454 CeedCallBackend(CeedFree(&is_node)); 455 456 // Compute transpose offsets and indices 457 const CeedInt size_offsets = num_nodes + 1; 458 459 CeedCallBackend(CeedCalloc(size_offsets, &t_offsets)); 460 CeedCallBackend(CeedMalloc(size_indices, &t_indices)); 461 // Count node multiplicity 462 for (CeedInt e = 0; e < num_elem; ++e) { 463 for (CeedInt i = 0; i < elem_size; ++i) ++t_offsets[ind_to_offset[indices[elem_size * e + i]] + 1]; 464 } 465 // Convert to running sum 466 for (CeedInt i = 1; i < size_offsets; ++i) t_offsets[i] += t_offsets[i - 1]; 467 // List all E-vec indices associated with L-vec node 468 for (CeedInt e = 0; e < num_elem; ++e) { 469 for (CeedInt i = 0; i < elem_size; ++i) { 470 const CeedInt lid = elem_size * e + i; 471 const CeedInt gid = indices[lid]; 472 473 t_indices[t_offsets[ind_to_offset[gid]]++] = lid; 474 } 475 } 476 // Reset running sum 477 for (int i = size_offsets - 1; i > 0; --i) t_offsets[i] = t_offsets[i - 1]; 478 t_offsets[0] = 0; 479 480 // Copy data to device 481 // -- L-vector indices 482 CeedCallHip(ceed, hipMalloc((void **)&impl->d_l_vec_indices, num_nodes * sizeof(CeedInt))); 483 CeedCallHip(ceed, hipMemcpy((CeedInt *)impl->d_l_vec_indices, l_vec_indices, num_nodes * sizeof(CeedInt), hipMemcpyHostToDevice)); 484 // -- Transpose offsets 485 CeedCallHip(ceed, hipMalloc((void **)&impl->d_t_offsets, size_offsets * sizeof(CeedInt))); 486 CeedCallHip(ceed, hipMemcpy((CeedInt *)impl->d_t_offsets, t_offsets, size_offsets * sizeof(CeedInt), hipMemcpyHostToDevice)); 487 // -- Transpose indices 488 CeedCallHip(ceed, hipMalloc((void **)&impl->d_t_indices, size_indices * sizeof(CeedInt))); 489 CeedCallHip(ceed, hipMemcpy((CeedInt *)impl->d_t_indices, t_indices, size_indices * sizeof(CeedInt), hipMemcpyHostToDevice)); 490 491 // Cleanup 492 CeedCallBackend(CeedFree(&ind_to_offset)); 493 CeedCallBackend(CeedFree(&l_vec_indices)); 494 CeedCallBackend(CeedFree(&t_offsets)); 495 CeedCallBackend(CeedFree(&t_indices)); 496 return CEED_ERROR_SUCCESS; 497 } 498 499 //------------------------------------------------------------------------------ 500 // Create restriction 501 //------------------------------------------------------------------------------ 502 int CeedElemRestrictionCreate_Hip(CeedMemType mem_type, CeedCopyMode copy_mode, const CeedInt *offsets, const bool *orients, 503 const CeedInt8 *curl_orients, CeedElemRestriction rstr) { 504 Ceed ceed, ceed_parent; 505 bool is_deterministic; 506 CeedInt num_elem, num_comp, elem_size; 507 CeedRestrictionType rstr_type; 508 CeedElemRestriction_Hip *impl; 509 510 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed)); 511 CeedCallBackend(CeedGetParent(ceed, &ceed_parent)); 512 CeedCallBackend(CeedIsDeterministic(ceed_parent, &is_deterministic)); 513 CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem)); 514 CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp)); 515 CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size)); 516 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type)); 517 // Use max number of points as elem size for AtPoints restrictions 518 if (rstr_type == CEED_RESTRICTION_POINTS) { 519 CeedInt max_points = 0; 520 521 for (CeedInt i = 0; i < num_elem; i++) { 522 max_points = CeedIntMax(max_points, offsets[i + 1] - offsets[i]); 523 } 524 elem_size = max_points; 525 } 526 const CeedInt size = num_elem * elem_size; 527 528 CeedCallBackend(CeedCalloc(1, &impl)); 529 impl->num_nodes = size; 530 CeedCallBackend(CeedElemRestrictionSetData(rstr, impl)); 531 532 // Set layouts 533 { 534 bool has_backend_strides; 535 CeedInt layout[3] = {1, size, elem_size}; 536 537 CeedCallBackend(CeedElemRestrictionSetELayout(rstr, layout)); 538 if (rstr_type == CEED_RESTRICTION_STRIDED) { 539 CeedCallBackend(CeedElemRestrictionHasBackendStrides(rstr, &has_backend_strides)); 540 if (has_backend_strides) { 541 CeedCallBackend(CeedElemRestrictionSetLLayout(rstr, layout)); 542 } 543 } 544 } 545 546 // Pad AtPoints indices 547 if (rstr_type == CEED_RESTRICTION_POINTS) { 548 CeedSize offsets_len = elem_size * num_elem, at_points_size = num_elem + 1; 549 CeedInt max_points = elem_size, *offsets_padded; 550 551 CeedCheck(mem_type == CEED_MEM_HOST, ceed, CEED_ERROR_BACKEND, "only MemType Host supported when creating AtPoints restriction"); 552 CeedCallBackend(CeedMalloc(offsets_len, &offsets_padded)); 553 for (CeedInt i = 0; i < num_elem; i++) { 554 CeedInt num_points = offsets[i + 1] - offsets[i]; 555 556 at_points_size += num_points; 557 // -- Copy all points in element 558 for (CeedInt j = 0; j < num_points; j++) { 559 offsets_padded[i * max_points + j] = offsets[offsets[i] + j] * num_comp; 560 } 561 // -- Replicate out last point in element 562 for (CeedInt j = num_points; j < max_points; j++) { 563 offsets_padded[i * max_points + j] = offsets[offsets[i] + num_points - 1] * num_comp; 564 } 565 } 566 CeedCallBackend(CeedSetHostCeedIntArray(offsets, copy_mode, at_points_size, &impl->h_offsets_at_points_owned, &impl->h_offsets_at_points_borrowed, 567 &impl->h_offsets_at_points)); 568 CeedCallHip(ceed, hipMalloc((void **)&impl->d_offsets_at_points_owned, at_points_size * sizeof(CeedInt))); 569 CeedCallHip(ceed, hipMemcpy((CeedInt **)impl->d_offsets_at_points_owned, impl->h_offsets_at_points, at_points_size * sizeof(CeedInt), 570 hipMemcpyHostToDevice)); 571 impl->d_offsets_at_points = (CeedInt *)impl->d_offsets_at_points_owned; 572 573 // -- Use padded offsets for the rest of the setup 574 offsets = (const CeedInt *)offsets_padded; 575 copy_mode = CEED_OWN_POINTER; 576 CeedCallBackend(CeedElemRestrictionSetAtPointsEVectorSize(rstr, at_points_size * num_comp)); 577 } 578 579 // Set up device offset/orientation arrays 580 if (rstr_type != CEED_RESTRICTION_STRIDED) { 581 switch (mem_type) { 582 case CEED_MEM_HOST: { 583 CeedCallBackend(CeedSetHostCeedIntArray(offsets, copy_mode, size, &impl->h_offsets_owned, &impl->h_offsets_borrowed, &impl->h_offsets)); 584 CeedCallHip(ceed, hipMalloc((void **)&impl->d_offsets_owned, size * sizeof(CeedInt))); 585 CeedCallHip(ceed, hipMemcpy((CeedInt **)impl->d_offsets_owned, impl->h_offsets, size * sizeof(CeedInt), hipMemcpyHostToDevice)); 586 impl->d_offsets = (CeedInt *)impl->d_offsets_owned; 587 if (is_deterministic) CeedCallBackend(CeedElemRestrictionOffset_Hip(rstr, elem_size, offsets)); 588 } break; 589 case CEED_MEM_DEVICE: { 590 CeedCallBackend(CeedSetDeviceCeedIntArray_Hip(ceed, offsets, copy_mode, size, &impl->d_offsets_owned, &impl->d_offsets_borrowed, 591 (const CeedInt **)&impl->d_offsets)); 592 CeedCallBackend(CeedMalloc(size, &impl->h_offsets_owned)); 593 CeedCallHip(ceed, hipMemcpy((CeedInt **)impl->h_offsets_owned, impl->d_offsets, size * sizeof(CeedInt), hipMemcpyDeviceToHost)); 594 impl->h_offsets = impl->h_offsets_owned; 595 if (is_deterministic) CeedCallBackend(CeedElemRestrictionOffset_Hip(rstr, elem_size, offsets)); 596 } break; 597 } 598 599 // Orientation data 600 if (rstr_type == CEED_RESTRICTION_ORIENTED) { 601 switch (mem_type) { 602 case CEED_MEM_HOST: { 603 CeedCallBackend(CeedSetHostBoolArray(orients, copy_mode, size, &impl->h_orients_owned, &impl->h_orients_borrowed, &impl->h_orients)); 604 CeedCallHip(ceed, hipMalloc((void **)&impl->d_orients_owned, size * sizeof(bool))); 605 CeedCallHip(ceed, hipMemcpy((bool *)impl->d_orients_owned, impl->h_orients, size * sizeof(bool), hipMemcpyHostToDevice)); 606 impl->d_orients = impl->d_orients_owned; 607 } break; 608 case CEED_MEM_DEVICE: { 609 CeedCallBackend(CeedSetDeviceBoolArray_Hip(ceed, orients, copy_mode, size, &impl->d_orients_owned, &impl->d_orients_borrowed, 610 (const bool **)&impl->d_orients)); 611 CeedCallBackend(CeedMalloc(size, &impl->h_orients_owned)); 612 CeedCallHip(ceed, hipMemcpy((bool *)impl->h_orients_owned, impl->d_orients, size * sizeof(bool), hipMemcpyDeviceToHost)); 613 impl->h_orients = impl->h_orients_owned; 614 } break; 615 } 616 } else if (rstr_type == CEED_RESTRICTION_CURL_ORIENTED) { 617 switch (mem_type) { 618 case CEED_MEM_HOST: { 619 CeedCallBackend(CeedSetHostCeedInt8Array(curl_orients, copy_mode, 3 * size, &impl->h_curl_orients_owned, &impl->h_curl_orients_borrowed, 620 &impl->h_curl_orients)); 621 CeedCallHip(ceed, hipMalloc((void **)&impl->d_curl_orients_owned, 3 * size * sizeof(CeedInt8))); 622 CeedCallHip(ceed, 623 hipMemcpy((CeedInt8 *)impl->d_curl_orients_owned, impl->h_curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyHostToDevice)); 624 impl->d_curl_orients = impl->d_curl_orients_owned; 625 } break; 626 case CEED_MEM_DEVICE: { 627 CeedCallBackend(CeedSetDeviceCeedInt8Array_Hip(ceed, curl_orients, copy_mode, 3 * size, &impl->d_curl_orients_owned, 628 &impl->d_curl_orients_borrowed, (const CeedInt8 **)&impl->d_curl_orients)); 629 CeedCallBackend(CeedMalloc(3 * size, &impl->h_curl_orients_owned)); 630 CeedCallHip(ceed, 631 hipMemcpy((CeedInt8 *)impl->h_curl_orients_owned, impl->d_curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyDeviceToHost)); 632 impl->h_curl_orients = impl->h_curl_orients_owned; 633 } break; 634 } 635 } 636 } 637 638 // Register backend functions 639 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "Apply", CeedElemRestrictionApply_Hip)); 640 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "ApplyUnsigned", CeedElemRestrictionApplyUnsigned_Hip)); 641 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "ApplyUnoriented", CeedElemRestrictionApplyUnoriented_Hip)); 642 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetOffsets", CeedElemRestrictionGetOffsets_Hip)); 643 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetOrientations", CeedElemRestrictionGetOrientations_Hip)); 644 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetCurlOrientations", CeedElemRestrictionGetCurlOrientations_Hip)); 645 if (rstr_type == CEED_RESTRICTION_POINTS) { 646 CeedCallBackend( 647 CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetAtPointsElementOffset", CeedElemRestrictionGetAtPointsElementOffset_Hip)); 648 } 649 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "Destroy", CeedElemRestrictionDestroy_Hip)); 650 return CEED_ERROR_SUCCESS; 651 } 652 653 //------------------------------------------------------------------------------ 654