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