From 18efd6555cf22c108d6d10a8ea56dcbaae1cb8df Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Thu, 21 Nov 2019 12:33:12 +0100 Subject: [PATCH] add GPU checking before compilation for rocm 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 9334b94b7cf94..168a8ba1757cb 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)