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 9334b94..861719f 100644 --- a/python/tvm/autotvm/measure/measure_methods.py +++ b/python/tvm/autotvm/measure/measure_methods.py @@ -227,7 +227,8 @@ class RPCRunner(Runner): 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 @@ class LocalRunner(RPCRunner): 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) -- 2.7.4