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-hip-compile.h" 9 10 #include <ceed.h> 11 #include <ceed/backend.h> 12 #include <ceed/jit-tools.h> 13 #include <stdarg.h> 14 #include <string.h> 15 #include <hip/hiprtc.h> 16 17 #include <sstream> 18 19 #include "ceed-hip-common.h" 20 21 #define CeedChk_hiprtc(ceed, x) \ 22 do { \ 23 hiprtcResult result = static_cast<hiprtcResult>(x); \ 24 if (result != HIPRTC_SUCCESS) return CeedError((ceed), CEED_ERROR_BACKEND, hiprtcGetErrorString(result)); \ 25 } while (0) 26 27 #define CeedCallHiprtc(ceed, ...) \ 28 do { \ 29 int ierr_q_ = __VA_ARGS__; \ 30 CeedChk_hiprtc(ceed, ierr_q_); \ 31 } while (0) 32 33 //------------------------------------------------------------------------------ 34 // Compile HIP kernel 35 //------------------------------------------------------------------------------ 36 int CeedCompile_Hip(Ceed ceed, const char *source, hipModule_t *module, const CeedInt num_defines, ...) { 37 size_t ptx_size; 38 char *ptx; 39 const int num_opts = 4; 40 CeedInt num_jit_source_dirs = 0, num_jit_defines = 0; 41 const char **opts; 42 int runtime_version; 43 hiprtcProgram prog; 44 struct hipDeviceProp_t prop; 45 Ceed_Hip *ceed_data; 46 47 hipFree(0); // Make sure a Context exists for hiprtc 48 49 std::ostringstream code; 50 51 // Add hip runtime include statement for generation if runtime < 40400000 (implies ROCm < 4.5) 52 CeedCallHip(ceed, hipRuntimeGetVersion(&runtime_version)); 53 if (runtime_version < 40400000) { 54 code << "\n#include <hip/hip_runtime.h>\n"; 55 } 56 // With ROCm 4.5, need to include these definitions specifically for hiprtc (but cannot include the runtime header) 57 else { 58 code << "#include <stddef.h>\n"; 59 code << "#define __forceinline__ inline __attribute__((always_inline))\n"; 60 code << "#define HIP_DYNAMIC_SHARED(type, var) extern __shared__ type var[];\n"; 61 } 62 63 // Kernel specific options, such as kernel constants 64 if (num_defines > 0) { 65 va_list args; 66 va_start(args, num_defines); 67 char *name; 68 int val; 69 70 for (int i = 0; i < num_defines; i++) { 71 name = va_arg(args, char *); 72 val = va_arg(args, int); 73 code << "#define " << name << " " << val << "\n"; 74 } 75 va_end(args); 76 } 77 78 // Standard libCEED definitions for HIP backends 79 code << "#include <ceed/jit-source/hip/hip-jit.h>\n\n"; 80 81 // Non-macro options 82 CeedCallBackend(CeedCalloc(num_opts, &opts)); 83 opts[0] = "-default-device"; 84 CeedCallBackend(CeedGetData(ceed, (void **)&ceed_data)); 85 CeedCallHip(ceed, hipGetDeviceProperties(&prop, ceed_data->device_id)); 86 std::string arch_arg = "--gpu-architecture=" + std::string(prop.gcnArchName); 87 opts[1] = arch_arg.c_str(); 88 opts[2] = "-munsafe-fp-atomics"; 89 opts[3] = "-DCEED_RUNNING_JIT_PASS=1"; 90 // Additional include dirs 91 { 92 const char **jit_source_dirs; 93 94 CeedCallBackend(CeedGetJitSourceRoots(ceed, &num_jit_source_dirs, &jit_source_dirs)); 95 CeedCallBackend(CeedRealloc(num_opts + num_jit_source_dirs, &opts)); 96 for (CeedInt i = 0; i < num_jit_source_dirs; i++) { 97 std::ostringstream include_dir_arg; 98 99 include_dir_arg << "-I" << jit_source_dirs[i]; 100 CeedCallBackend(CeedStringAllocCopy(include_dir_arg.str().c_str(), (char **)&opts[num_opts + i])); 101 } 102 CeedCallBackend(CeedRestoreJitSourceRoots(ceed, &jit_source_dirs)); 103 } 104 // User defines 105 { 106 const char **jit_defines; 107 108 CeedCallBackend(CeedGetJitDefines(ceed, &num_jit_defines, &jit_defines)); 109 CeedCallBackend(CeedRealloc(num_opts + num_jit_source_dirs + num_jit_defines, &opts)); 110 for (CeedInt i = 0; i < num_jit_defines; i++) { 111 std::ostringstream define_arg; 112 113 define_arg << "-D" << jit_defines[i]; 114 CeedCallBackend(CeedStringAllocCopy(define_arg.str().c_str(), (char **)&opts[num_opts + num_jit_source_dirs + i])); 115 } 116 CeedCallBackend(CeedRestoreJitDefines(ceed, &jit_defines)); 117 } 118 119 // Add string source argument provided in call 120 code << source; 121 122 // Create Program 123 CeedCallHiprtc(ceed, hiprtcCreateProgram(&prog, code.str().c_str(), NULL, 0, NULL, NULL)); 124 125 // Compile kernel 126 CeedDebug256(ceed, CEED_DEBUG_COLOR_ERROR, "---------- ATTEMPTING TO COMPILE JIT SOURCE ----------\n"); 127 CeedDebug(ceed, "Source:\n%s\n", code.str().c_str()); 128 CeedDebug256(ceed, CEED_DEBUG_COLOR_ERROR, "---------- END OF JIT SOURCE ----------\n"); 129 hiprtcResult result = hiprtcCompileProgram(prog, num_opts + num_jit_source_dirs + num_jit_defines, opts); 130 131 for (CeedInt i = 0; i < num_jit_source_dirs; i++) { 132 CeedCallBackend(CeedFree(&opts[num_opts + i])); 133 } 134 for (CeedInt i = 0; i < num_jit_defines; i++) { 135 CeedCallBackend(CeedFree(&opts[num_opts + num_jit_source_dirs + i])); 136 } 137 CeedCallBackend(CeedFree(&opts)); 138 if (result != HIPRTC_SUCCESS) { 139 size_t log_size; 140 char *log; 141 142 CeedChk_hiprtc(ceed, hiprtcGetProgramLogSize(prog, &log_size)); 143 CeedCallBackend(CeedMalloc(log_size, &log)); 144 CeedCallHiprtc(ceed, hiprtcGetProgramLog(prog, log)); 145 return CeedError(ceed, CEED_ERROR_BACKEND, "%s\n%s", hiprtcGetErrorString(result), log); 146 } 147 148 CeedCallHiprtc(ceed, hiprtcGetCodeSize(prog, &ptx_size)); 149 CeedCallBackend(CeedMalloc(ptx_size, &ptx)); 150 CeedCallHiprtc(ceed, hiprtcGetCode(prog, ptx)); 151 CeedCallHiprtc(ceed, hiprtcDestroyProgram(&prog)); 152 153 CeedCallHip(ceed, hipModuleLoadData(module, ptx)); 154 CeedCallBackend(CeedFree(&ptx)); 155 return CEED_ERROR_SUCCESS; 156 } 157 158 //------------------------------------------------------------------------------ 159 // Get HIP kernel 160 //------------------------------------------------------------------------------ 161 int CeedGetKernel_Hip(Ceed ceed, hipModule_t module, const char *name, hipFunction_t *kernel) { 162 CeedCallHip(ceed, hipModuleGetFunction(kernel, module, name)); 163 return CEED_ERROR_SUCCESS; 164 } 165 166 //------------------------------------------------------------------------------ 167 // Run HIP kernel 168 //------------------------------------------------------------------------------ 169 int CeedRunKernel_Hip(Ceed ceed, hipFunction_t kernel, const int grid_size, const int block_size, void **args) { 170 CeedCallHip(ceed, hipModuleLaunchKernel(kernel, grid_size, 1, 1, block_size, 1, 1, 0, NULL, args, NULL)); 171 return CEED_ERROR_SUCCESS; 172 } 173 174 //------------------------------------------------------------------------------ 175 // Run HIP kernel for spatial dimension 176 //------------------------------------------------------------------------------ 177 int CeedRunKernelDim_Hip(Ceed ceed, hipFunction_t kernel, const int grid_size, const int block_size_x, const int block_size_y, const int block_size_z, 178 void **args) { 179 CeedCallHip(ceed, hipModuleLaunchKernel(kernel, grid_size, 1, 1, block_size_x, block_size_y, block_size_z, 0, NULL, args, NULL)); 180 return CEED_ERROR_SUCCESS; 181 } 182 183 //------------------------------------------------------------------------------ 184 // Run HIP kernel for spatial dimension with shared memory 185 //------------------------------------------------------------------------------ 186 int CeedRunKernelDimShared_Hip(Ceed ceed, hipFunction_t kernel, const int grid_size, const int block_size_x, const int block_size_y, 187 const int block_size_z, const int shared_mem_size, void **args) { 188 CeedCallHip(ceed, hipModuleLaunchKernel(kernel, grid_size, 1, 1, block_size_x, block_size_y, block_size_z, shared_mem_size, NULL, args, NULL)); 189 return CEED_ERROR_SUCCESS; 190 } 191 192 //------------------------------------------------------------------------------ 193