[DRIVER] Fixed some issue with how ptxas is used (#399)

Now using tmpnam and properly deleting temporaries when an exception is raised
This commit is contained in:
Philippe Tillet
2021-12-21 14:31:51 -08:00
committed by GitHub
parent 39d4bfed83
commit 2509124dd0

View File

@@ -178,16 +178,14 @@ std::string ptx_to_cubin(const std::string& ptx, int cc) {
" but a working version could not be found."); " but a working version could not be found.");
std::string ptxas = working_ptxas.front(); std::string ptxas = working_ptxas.front();
// compile ptx with ptxas // compile ptx with ptxas
char _fsrc[] = "/tmp/triton_k_XXXXXX"; char _fsrc[L_tmpnam];
char _flog[] = "/tmp/triton_l_XXXXXX"; char _flog[L_tmpnam];
mkstemp(_fsrc); std::string fsrc = std::tmpnam(_fsrc);
mkstemp(_flog); std::string flog = std::tmpnam(_flog);
std::string fsrc = _fsrc;
std::string flog = _flog;
std::string fbin = fsrc + ".o"; std::string fbin = fsrc + ".o";
const char* _fbin = fbin.c_str(); const char* _fbin = fbin.c_str();
std::ofstream ofs(fsrc); std::ofstream ofs(fsrc);
ofs << ptx; ofs << ptx << std::endl;
ofs.close(); ofs.close();
std::string cmd; std::string cmd;
int err; int err;
@@ -196,16 +194,18 @@ std::string ptx_to_cubin(const std::string& ptx, int cc) {
if(err != 0){ if(err != 0){
std::ifstream _log(_flog); std::ifstream _log(_flog);
std::string log(std::istreambuf_iterator<char>(_log), {}); std::string log(std::istreambuf_iterator<char>(_log), {});
unlink(_fsrc);
unlink(_flog);
throw std::runtime_error("Internal Triton PTX codegen error: \n" + log); throw std::runtime_error("Internal Triton PTX codegen error: \n" + log);
} }
CUmodule ret; CUmodule ret;
std::ifstream _cubin(_fbin, std::ios::binary ); std::ifstream _cubin(_fbin, std::ios::binary );
std::string cubin(std::istreambuf_iterator<char>(_cubin), {}); std::string cubin(std::istreambuf_iterator<char>(_cubin), {});
_cubin.close(); _cubin.close();
dispatch::cuModuleLoadData(&ret, cubin.c_str());
unlink(_fsrc); unlink(_fsrc);
unlink(_flog); unlink(_flog);
unlink(_fbin); unlink(_fbin);
dispatch::cuModuleLoadData(&ret, cubin.c_str());
return cubin; return cubin;
} }