From 786d7998e1afe0a9d5a9e0091a4509afc748d798 Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Thu, 21 Nov 2019 15:40:29 +0100 Subject: [PATCH] add GPU checking before compilation for rocm (#4394) 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. --- python/tvm/autotvm/measure/measure_methods.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/autotvm/measure/measure_methods.py b/python/tvm/autotvm/measure/measure_methods.py index 9334b94b7cf9..861719fffe53 100644 --- a/python/tvm/autotvm/measure/measure_methods.py +++ b/python/tvm/autotvm/measure/measure_methods.py @@ -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 @@ -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)