Skip to content

Commit

Permalink
add GPU checking before compilation for rocm
Browse files Browse the repository at this point in the history
Previously, we would rely on the later phases to error out
(often for using too much shared memory). This enables the
checks on the IR that already exist for CUDA and OpenCL also
for ROCm.
  • Loading branch information
t-vi committed Nov 21, 2019
1 parent 464ebb1 commit 18efd65
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions python/tvm/autotvm/measure/measure_methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,8 @@ def set_task(self, task):

def get_build_kwargs(self):
kwargs = {}
if 'cuda' in self.task.target.keys or 'opencl' in self.task.target.keys:
if 'cuda' in self.task.target.keys or 'opencl' in self.task.target.keys or \
'rocm in self.task.target.keys':
remote = request_remote(self.key, self.host, self.port)
ctx = remote.context(str(self.task.target), 0)
max_dims = ctx.max_thread_dimensions
Expand Down Expand Up @@ -345,7 +346,6 @@ def set_task(self, task):
def _build_func_common(measure_input, check_gpu=None, cuda_arch=None, build_option=None):
"""Common part for building a configuration"""
target, task, config = measure_input

with target:
s, args = task.instantiate(config)

Expand Down

0 comments on commit 18efd65

Please sign in to comment.