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