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