From eae6588a2ac337ee6bac1427d9c3c58503ff01d0 Mon Sep 17 00:00:00 2001 From: Henry Lu Date: Thu, 24 Oct 2019 13:32:05 +0800 Subject: [PATCH 1/2] Fix example code in comment of tvm.build_module.build() --- python/tvm/build_module.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index 4cb09931616e..c05ebb96493c 100644 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -568,10 +568,11 @@ def build(inputs, B = tvm.placeholder((n,), name='B') C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') s1 = tvm.create_schedule(C.op) - s2 = topi.cpp.cuda.schedule_injective("cuda", [C]) - f1 = tvm.lower(s1, [A, B, C], name="test_add1") - f2 = tvm.lower(s2, [A, B, C], name="test_add2") - m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm") + with tvm.target.cuda() as cuda_tgt: + s2 = topi.cpp.cuda.schedule_injective(cuda_tgt, [C]) + f1 = tvm.lower(s1, [A, B, C], name="test_add1") + f2 = tvm.lower(s2, [A, B, C], name="test_add2") + m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm") Note ---- From da239faf15512a73ff3526f66700e97492f018e9 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 24 Oct 2019 09:55:48 -0700 Subject: [PATCH 2/2] Update build_module.py --- python/tvm/build_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index c05ebb96493c..479bdcead7c5 100644 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -569,7 +569,7 @@ def build(inputs, C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') s1 = tvm.create_schedule(C.op) with tvm.target.cuda() as cuda_tgt: - s2 = topi.cpp.cuda.schedule_injective(cuda_tgt, [C]) + s2 = topi.cuda.schedule_injective(cuda_tgt, [C]) f1 = tvm.lower(s1, [A, B, C], name="test_add1") f2 = tvm.lower(s2, [A, B, C], name="test_add2") m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm")