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

[PYTHON] Restored compatibility with powerpc

Browse Source
This commit is contained in:
Philippe Tillet
2020-05-09 23:46:25 -04:00
committed by Philippe Tillet
parent 1426b103e9
commit f35b9100e2
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/triton/kernel.py
View File

@@ -19,6 +19,7 @@ import triton._C.libtriton as libtriton
import os import os
import time import time
import torch.utils.cpp_extension import torch.utils.cpp_extension
import platform
def _cvt_to_def_str(obj): def _cvt_to_def_str(obj):
@@ -85,10 +86,12 @@ def _make_framework_op(arg_types):
ccdir = os.path.join(libtriton.__file__, os.path.pardir) ccdir = os.path.join(libtriton.__file__, os.path.pardir)
ccdir = os.path.realpath(ccdir) ccdir = os.path.realpath(ccdir)
print('[TRITON] Compiling op...') print('[TRITON] Compiling op...')
machine = platform.machine()
extra_cflags = ['-std=gnu++11'] if machine == 'ppc64le' else None
lib = torch.utils.cpp_extension.load(name, cpp, lib = torch.utils.cpp_extension.load(name, cpp,
extra_ldflags = [f'-L{ccdir}', '-ltriton'], extra_ldflags = [f'-L{ccdir}', '-ltriton'],
extra_include_paths = [os.path.join(ccdir, 'include')], extra_include_paths = [os.path.join(ccdir, 'include')],
#extra_cflags = [f'-std=gnu++11'], extra_cflags = extra_cflags,
build_directory = root, build_directory = root,
is_python_module = False, is_python_module = False,
with_cuda = True) with_cuda = True)
@@ -146,4 +149,4 @@ class kernel:
# launch # launch
self.fw_op(self.op_id, device, bench, bench_id, *args) self.fw_op(self.op_id, device, bench, bench_id, *args)
if bench > 0: if bench > 0:
return libtriton.retrieve_scalar(bench_id) return libtriton.retrieve_scalar(bench_id)