Mark Shields
c8a6089073
* This continues the work outlined in the RFC https://discuss.tvm.apache.org/t/rfc-relay-tecompiler-rewrite-existing-compile-engine-to-match-updated-compiler-flow/9233 This gets about halfway there for the Interpreter: * Remove direct access to TECompiler from interpreter, and instead call tec::LowerTEExpr when 'preparing' a module and expression for evaluation. * Make clear there's no phase distinction between create_interpreter and evaluate on the Python side -- both must be prepared together as a single IRModule. * But in return make sure the result of evaluate on the Python side is a packed func ready to directly apply 'simple' arguments to an already interpreted closure. * The interpreter builds and caches primitive TIR functions (and their corresponding dynamic shape functions) as packed funcs as they are encountered. * Cleanup uses of interpreter for constant folding on the C++ side. Future work: * Fold LoweredModule into IRModule so tec::LowerTEExpr is just another pass. * Get rid of the implicit caching of lowered functions in TECompiler. * Make calling convention from Relay to TIR explicit, and remove all the function attribute hackery currently needed so the interpreter can correctly invoke lowered functions as it encounters them. * Make TECompiler private. Though could do this now it will make migrating the VM and AOT uses of CompilerEngine harder. Force a gc between sphinx-gallery items to reclaim GPU memory. (#8722) GPU memory is only released once the PackedFunc for evaling the model is gced by Python. In CI we're noticing intermittent 'CUDA: Out of memory' failures while processing the tutorials, and tracing showed there was no gc happening between items. Not confident this will solve the problem but worth a try. * Get rid of logs spam.
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
from tvm.relay import var, const, create_executor
|
|
from tvm.relay.op import debug
|
|
|
|
|
|
_test_debug_hit = False
|
|
|
|
|
|
def test_debug():
|
|
global _test_debug_hit
|
|
x = var("x", shape=(), dtype="int32")
|
|
_test_debug_hit = False
|
|
|
|
def did_exec(x):
|
|
global _test_debug_hit
|
|
_test_debug_hit = True
|
|
|
|
prog = debug(x, debug_func=did_exec)
|
|
result = create_executor().evaluate(prog, {x: const(1, "int32")})
|
|
assert _test_debug_hit
|
|
assert result.numpy() == 1
|
|
|
|
|
|
def test_debug_with_expr():
|
|
global _test_debug_hit
|
|
_test_debug_hit = False
|
|
x = var("x", shape=(), dtype="int32")
|
|
_test_debug_hit = False
|
|
|
|
def did_exec(x):
|
|
global _test_debug_hit
|
|
_test_debug_hit = True
|
|
|
|
prog = debug(x + x * x, debug_func=did_exec)
|
|
result = create_executor().evaluate(prog, {x: const(2, "int32")})
|
|
assert _test_debug_hit
|
|
assert result.numpy() == 6
|