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