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