gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[BUILD] Fix Warnings and Enable Warnings as Errors (#794)

Browse Source
This commit is contained in:
Ian Bearman
2022-10-28 12:36:09 -07:00
committed by GitHub
parent ac0f6793cc
commit f2106d0aa2
20 changed files with 205 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
python/setup.py
View File

@@ -1,11 +1,12 @@
import distutils
import distutils.spawn
import itertools
import os
import platform
import re
import shutil
import subprocess
import sys
import sysconfig
import tarfile
import tempfile
import urllib.request
@@ -16,6 +17,74 @@ from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
# Logic from https://github.com/Kitware/CMake/blob/master/Modules/FindPythonLibs.cmake
# Code from https://stackoverflow.com/questions/47423246
def get_python_library():
"""Get path to the python library associated with the current python
interpreter."""
# determine direct path to libpython
python_version = sysconfig.get_python_version()
python_library = sysconfig.get_config_var('LIBRARY')
# if static (or nonexistent), try to find a suitable dynamic libpython
if (python_library is None or os.path.splitext(python_library)[1][-2:] == '.a'):
candidate_lib_prefixes = ['', 'lib']
candidate_extensions = ['.lib', '.so', '.a']
if sysconfig.get_config_var('WITH_DYLD'):
candidate_extensions.insert(0, '.dylib')
candidate_versions = [python_version]
if python_version:
candidate_versions.append('')
candidate_versions.insert(
0, "".join(python_version.split(".")[:2]))
abiflags = getattr(sys, 'abiflags', '')
candidate_abiflags = [abiflags]
if abiflags:
candidate_abiflags.append('')
# Ensure the value injected by virtualenv is
# returned on windows.
# Because calling `sysconfig.get_config_var('multiarchsubdir')`
# returns an empty string on Linux, `du_sysconfig` is only used to
# get the value of `LIBDIR`.
libdir = distutils.sysconfig.get_config_var('LIBDIR')
if sysconfig.get_config_var('MULTIARCH'):
masd = sysconfig.get_config_var('multiarchsubdir')
if masd:
if masd.startswith(os.sep):
masd = masd[len(os.sep):]
libdir = os.path.join(libdir, masd)
if libdir is None:
libdir = os.path.abspath(os.path.join(
sysconfig.get_config_var('LIBDEST'), "..", "libs"))
candidates = (
os.path.join(
libdir,
''.join((pre, 'python', ver, abi, ext))
)
for (pre, ext, ver, abi) in itertools.product(
candidate_lib_prefixes,
candidate_extensions,
candidate_versions,
candidate_abiflags
)
)
for candidate in candidates:
if os.path.exists(candidate):
# we found a (likely alternate) libpython
python_library = candidate
break
return python_library
# Taken from https://github.com/pytorch/pytorch/blob/master/tools/setup_helpers/env.py
def check_env_flag(name: str, default: str = "") -> bool:
return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
@@ -136,14 +205,19 @@ class CMakeBuild(build_ext):
if not os.path.exists(llvm_build_dir):
os.makedirs(llvm_build_dir)
# python directories
python_include_dirs = [distutils.sysconfig.get_python_inc()] + ['/usr/local/cuda/include']
python_include_dir = distutils.sysconfig.get_python_inc()
python_link_dir = distutils.sysconfig.get_python_lib()
python_library = get_python_library()
cmake_args = [
"-DLLVM_ENABLE_WERROR=ON",
"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir,
"-DTRITON_BUILD_TUTORIALS=OFF",
"-DTRITON_BUILD_PYTHON_MODULE=ON",
# '-DPYTHON_EXECUTABLE=' + sys.executable,
# '-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON',
"-DPYTHON_INCLUDE_DIRS=" + ";".join(python_include_dirs),
"-DPYTHON_INCLUDE_DIRS=" + python_include_dir,
"-DPYTHON_LINK_DIRS=" + python_link_dir,
"-DPYTHON_LIBRARIES=" + python_library,
"-DLLVM_EXTERNAL_LIT=" + lit_dir
] + thirdparty_cmake_args

36
python/src/triton.cc
View File

@@ -26,6 +26,7 @@
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/IRReader/IRReader.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/SourceMgr.h"
@@ -1301,39 +1302,36 @@ void init_triton_translation(py::module &m) {
py::gil_scoped_release allow_threads;
// compile ptx with ptxas
char _fsrc[L_tmpnam];
char _flog[L_tmpnam];
std::tmpnam(_fsrc);
std::tmpnam(_flog);
std::string fsrc = _fsrc;
std::string flog = _flog;
std::string fbin = fsrc + ".o";
llvm::SmallString<64> fsrc;
llvm::SmallString<64> flog;
llvm::sys::fs::createTemporaryFile("compile-ptx-src", "", fsrc);
llvm::sys::fs::createTemporaryFile("compile-ptx-log", "", flog);
std::string fbin = std::string(fsrc) + ".o";
llvm::FileRemover srcRemover(fsrc);
llvm::FileRemover logRemover(flog);
llvm::FileRemover binRemover(fbin);
const char *_fsrc = fsrc.c_str();
const char *_flog = flog.c_str();
const char *_fbin = fbin.c_str();
std::ofstream ofs(fsrc);
std::ofstream ofs(_fsrc);
ofs << ptxCode << std::endl;
ofs.close();
std::string cmd;
int err;
cmd = ptxasPath + " -v --gpu-name=sm_" + std::to_string(capability) +
" " + fsrc + " -o " + fsrc + ".o 2> " + flog;
" " + _fsrc + " -o " + _fsrc + ".o 2> " + _flog;
err = system(cmd.c_str());
if (err != 0) {
std::ifstream _log(_flog);
std::string log(std::istreambuf_iterator<char>(_log), {});
unlink(_fsrc);
unlink(_flog);
throw std::runtime_error("Internal Triton PTX codegen error: \n" +
log);
}
std::ifstream _cubin(_fbin, std::ios::binary);
std::string cubin(std::istreambuf_iterator<char>(_cubin), {});
_cubin.close();
unlink(_fsrc);
unlink(_flog);
unlink(_fbin);
py::bytes bytes(cubin);
return bytes;
return std::move(bytes);
});
m.def("add_external_libs",
@@ -1345,8 +1343,8 @@ void init_triton_translation(py::module &m) {
void init_triton(py::module &m) {
py::module subm = m.def_submodule("triton");
// init_triton_codegen(std::move(subm.def_submodule("code_gen")));
init_triton_runtime(std::move(subm.def_submodule("runtime")));
init_triton_ir(std::move(subm.def_submodule("ir")));
// init_triton_codegen(subm.def_submodule("code_gen"));
init_triton_runtime(subm.def_submodule("runtime"));
init_triton_ir(subm.def_submodule("ir"));
init_triton_translation(subm);
}