From 70581364771e2415b37b202a9fa6a937f275cfc6 Mon Sep 17 00:00:00 2001 From: Mei Ye Date: Tue, 1 Feb 2022 01:39:50 +0000 Subject: [PATCH 1/3] [relay] Fix stack overflow in device_planner observed on windows due to recursive function calls. --- src/relay/transforms/device_planner.cc | 108 +++++++++++-------- tests/python/relay/test_pass_plan_devices.py | 42 ++++++++ 2 files changed, 106 insertions(+), 44 deletions(-) diff --git a/src/relay/transforms/device_planner.cc b/src/relay/transforms/device_planner.cc index c7ca2227fd90..7ed312e3aeca 100644 --- a/src/relay/transforms/device_planner.cc +++ b/src/relay/transforms/device_planner.cc @@ -489,51 +489,71 @@ class DeviceAnalyzer : public ExprVisitor { void VisitExpr_(const CallNode* call_node) final { auto call = GetRef(call_node); - - // We don't care if the call is in pre- or post-lowered form. - auto vanilla_call = GetAnyCall(call_node); - - // Find the higher-order domain for the callee. See DomainForCallee for the special rules - // for primitives. - VisitExpr(vanilla_call->op); - auto func_domain = domains_->DomainForCallee(call); // higher-order - - // Build the domain for the function implied by its arguments and call context. - ICHECK_EQ(func_domain->function_arity(), vanilla_call->args.size()) << PrettyPrint(call); - std::vector args_and_result_domains; - args_and_result_domains.reserve(vanilla_call->args.size() + 1); - for (const auto& arg : vanilla_call->args) { - args_and_result_domains.emplace_back(domains_->DomainFor(arg)); - VisitExpr(arg); - } - args_and_result_domains.emplace_back(domains_->DomainFor(call)); - auto implied_domain = - domains_->MakeHigherOrderDomain(std::move(args_and_result_domains)); // higher-order - - VLOG(2) << "initial call function domain:" << std::endl - << domains_->ToString(func_domain) << std::endl - << "and implied domain:" << std::endl - << domains_->ToString(implied_domain) << std::endl - << "for call:" << std::endl - << PrettyPrint(call); - - // The above must match. - if (domains_->UnifyOrNull(func_domain, implied_domain) == nullptr) { // higher-order - // TODO(mbs): Proper diagnostics. - LOG(FATAL) - << "Function parameters and result VirtualDevices do not match those of call. Call:" - << std::endl - << PrettyPrint(call) << std::endl - << "with function virtual devices:" << std::endl - << domains_->ToString(func_domain) << std::endl - << "and implied call virtual devices:" << std::endl - << domains_->ToString(implied_domain); + std::stack stack; + std::unordered_map visited; + stack.push(call); + + while (stack.size() > 0) { + auto cur_expr = stack.top(); + auto* cur_node = cur_expr.as(); + if (cur_node != nullptr) { + bool is_visited = (visited.find(cur_node) != visited.end()); + auto vanilla_call = GetAnyCall(cur_node); + if (is_visited) { + auto call = GetRef(cur_node); + auto func_domain = domains_->DomainForCallee(call); // higher-order + ICHECK_EQ(func_domain->function_arity(), vanilla_call->args.size()) << PrettyPrint(call); + std::vector args_and_result_domains; + args_and_result_domains.reserve(vanilla_call->args.size() + 1); + for (const auto& arg : vanilla_call->args) { + args_and_result_domains.emplace_back(domains_->DomainFor(arg)); + } + args_and_result_domains.emplace_back(domains_->DomainFor(call)); + auto implied_domain = + domains_->MakeHigherOrderDomain(std::move(args_and_result_domains)); // higher-order + + VLOG(2) << "initial call function domain:" << std::endl + << domains_->ToString(func_domain) << std::endl + << "and implied domain:" << std::endl + << domains_->ToString(implied_domain) << std::endl + << "for call:" << std::endl + << PrettyPrint(call); + + // The above must match. + if (domains_->UnifyOrNull(func_domain, implied_domain) == nullptr) { // higher-order + // TODO(mbs): Proper diagnostics. + LOG(FATAL) + << "Function parameters and result SEScopes do not match those of call. Call:" + << std::endl + << PrettyPrint(call) << std::endl + << "with function scopes:" << std::endl + << domains_->ToString(func_domain) << std::endl + << "and implied call scopes:" << std::endl + << domains_->ToString(implied_domain); + } + + VLOG(2) << "final call function domain:" << std::endl + << domains_->ToString(func_domain) << std::endl + << "for call:" << std::endl + << PrettyPrint(call); + + stack.pop(); + } else { + VisitExpr(vanilla_call->op); + visited.emplace(cur_node, true); + + for (auto it = vanilla_call->args.rbegin(); it != vanilla_call->args.rend(); ++it) { + auto* arg_node = (*it).as(); + if ((arg_node == nullptr) || (visited.find(arg_node) == visited.end())) { + stack.push(*it); + } + } + } + } else { + VisitExpr(cur_expr); + stack.pop(); + } } - - VLOG(2) << "final call function domain:" << std::endl - << domains_->ToString(func_domain) << std::endl - << "for call:" << std::endl - << PrettyPrint(call); } void VisitExpr_(const LetNode* let_node) final { diff --git a/tests/python/relay/test_pass_plan_devices.py b/tests/python/relay/test_pass_plan_devices.py index 6232f9c4a96a..cbed1adb19e2 100644 --- a/tests/python/relay/test_pass_plan_devices.py +++ b/tests/python/relay/test_pass_plan_devices.py @@ -26,6 +26,7 @@ from tvm.script import tir as T import tvm.testing import numpy as np +import os HOST_DEVICE = tvm.device("cpu") HOST_TARGET = tvm.target.Target("llvm") @@ -1691,6 +1692,47 @@ def @main(%x : Tensor[(128, 128), float32], exercise(input(), expected(), None, None) +def test_stack_overflow(): + metatable = {"VirtualDevice": [CPU, GPU]} + + # Everything defaults to GPU + def input(): + tmp = "test_stack_overflow_input.txt" + with open(tmp, "w") as f: + f.write( + """ + #[version = "0.0.5"] + def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32], + %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32]) { + %0 = add(%a, %b); + %1 = add(%c, %d); + """ + ) + end = 1555 + for i in range(2, end): + s1 = "\n\t" + "%" + str(i) + " = add(%" + str(i - 1) + ", %" + str(i - 2) + ");" + f.write(s1) + f.write("\n\t" + "add(%" + str(end - 1) + ", %" + str(end - 2) + ")") + f.write("\n\t}") + + with open(tmp) as f: + mod = f.read() + f.close() + + os.remove(tmp) + + return tvm.parser.parse( + mod, + "from_string", + None, + metatable, + ) + config = tvm.target.make_compilation_config(CTXT, TARGETS, HOST_TARGET) + actual_mod = relay.transform.InferType()(in_mod) + actual_mod = relay.transform.PlanDevices(config)(actual_mod) + relay.transform.InferType()(actual_mod) + + if __name__ == "__main__": import sys import pytest From 5a935b5cadb0a4f518fe9e487a62b770265bd086 Mon Sep 17 00:00:00 2001 From: Mei Ye Date: Mon, 7 Feb 2022 23:24:22 +0000 Subject: [PATCH 2/3] Revert "[relay] Fix stack overflow in device_planner observed on windows due to recursive function calls." This reverts commit 70581364771e2415b37b202a9fa6a937f275cfc6. --- src/relay/transforms/device_planner.cc | 108 ++++++++----------- tests/python/relay/test_pass_plan_devices.py | 42 -------- 2 files changed, 44 insertions(+), 106 deletions(-) diff --git a/src/relay/transforms/device_planner.cc b/src/relay/transforms/device_planner.cc index 7ed312e3aeca..c7ca2227fd90 100644 --- a/src/relay/transforms/device_planner.cc +++ b/src/relay/transforms/device_planner.cc @@ -489,71 +489,51 @@ class DeviceAnalyzer : public ExprVisitor { void VisitExpr_(const CallNode* call_node) final { auto call = GetRef(call_node); - std::stack stack; - std::unordered_map visited; - stack.push(call); - - while (stack.size() > 0) { - auto cur_expr = stack.top(); - auto* cur_node = cur_expr.as(); - if (cur_node != nullptr) { - bool is_visited = (visited.find(cur_node) != visited.end()); - auto vanilla_call = GetAnyCall(cur_node); - if (is_visited) { - auto call = GetRef(cur_node); - auto func_domain = domains_->DomainForCallee(call); // higher-order - ICHECK_EQ(func_domain->function_arity(), vanilla_call->args.size()) << PrettyPrint(call); - std::vector args_and_result_domains; - args_and_result_domains.reserve(vanilla_call->args.size() + 1); - for (const auto& arg : vanilla_call->args) { - args_and_result_domains.emplace_back(domains_->DomainFor(arg)); - } - args_and_result_domains.emplace_back(domains_->DomainFor(call)); - auto implied_domain = - domains_->MakeHigherOrderDomain(std::move(args_and_result_domains)); // higher-order - - VLOG(2) << "initial call function domain:" << std::endl - << domains_->ToString(func_domain) << std::endl - << "and implied domain:" << std::endl - << domains_->ToString(implied_domain) << std::endl - << "for call:" << std::endl - << PrettyPrint(call); - - // The above must match. - if (domains_->UnifyOrNull(func_domain, implied_domain) == nullptr) { // higher-order - // TODO(mbs): Proper diagnostics. - LOG(FATAL) - << "Function parameters and result SEScopes do not match those of call. Call:" - << std::endl - << PrettyPrint(call) << std::endl - << "with function scopes:" << std::endl - << domains_->ToString(func_domain) << std::endl - << "and implied call scopes:" << std::endl - << domains_->ToString(implied_domain); - } - - VLOG(2) << "final call function domain:" << std::endl - << domains_->ToString(func_domain) << std::endl - << "for call:" << std::endl - << PrettyPrint(call); - - stack.pop(); - } else { - VisitExpr(vanilla_call->op); - visited.emplace(cur_node, true); - - for (auto it = vanilla_call->args.rbegin(); it != vanilla_call->args.rend(); ++it) { - auto* arg_node = (*it).as(); - if ((arg_node == nullptr) || (visited.find(arg_node) == visited.end())) { - stack.push(*it); - } - } - } - } else { - VisitExpr(cur_expr); - stack.pop(); - } + + // We don't care if the call is in pre- or post-lowered form. + auto vanilla_call = GetAnyCall(call_node); + + // Find the higher-order domain for the callee. See DomainForCallee for the special rules + // for primitives. + VisitExpr(vanilla_call->op); + auto func_domain = domains_->DomainForCallee(call); // higher-order + + // Build the domain for the function implied by its arguments and call context. + ICHECK_EQ(func_domain->function_arity(), vanilla_call->args.size()) << PrettyPrint(call); + std::vector args_and_result_domains; + args_and_result_domains.reserve(vanilla_call->args.size() + 1); + for (const auto& arg : vanilla_call->args) { + args_and_result_domains.emplace_back(domains_->DomainFor(arg)); + VisitExpr(arg); } + args_and_result_domains.emplace_back(domains_->DomainFor(call)); + auto implied_domain = + domains_->MakeHigherOrderDomain(std::move(args_and_result_domains)); // higher-order + + VLOG(2) << "initial call function domain:" << std::endl + << domains_->ToString(func_domain) << std::endl + << "and implied domain:" << std::endl + << domains_->ToString(implied_domain) << std::endl + << "for call:" << std::endl + << PrettyPrint(call); + + // The above must match. + if (domains_->UnifyOrNull(func_domain, implied_domain) == nullptr) { // higher-order + // TODO(mbs): Proper diagnostics. + LOG(FATAL) + << "Function parameters and result VirtualDevices do not match those of call. Call:" + << std::endl + << PrettyPrint(call) << std::endl + << "with function virtual devices:" << std::endl + << domains_->ToString(func_domain) << std::endl + << "and implied call virtual devices:" << std::endl + << domains_->ToString(implied_domain); + } + + VLOG(2) << "final call function domain:" << std::endl + << domains_->ToString(func_domain) << std::endl + << "for call:" << std::endl + << PrettyPrint(call); } void VisitExpr_(const LetNode* let_node) final { diff --git a/tests/python/relay/test_pass_plan_devices.py b/tests/python/relay/test_pass_plan_devices.py index cbed1adb19e2..6232f9c4a96a 100644 --- a/tests/python/relay/test_pass_plan_devices.py +++ b/tests/python/relay/test_pass_plan_devices.py @@ -26,7 +26,6 @@ from tvm.script import tir as T import tvm.testing import numpy as np -import os HOST_DEVICE = tvm.device("cpu") HOST_TARGET = tvm.target.Target("llvm") @@ -1692,47 +1691,6 @@ def @main(%x : Tensor[(128, 128), float32], exercise(input(), expected(), None, None) -def test_stack_overflow(): - metatable = {"VirtualDevice": [CPU, GPU]} - - # Everything defaults to GPU - def input(): - tmp = "test_stack_overflow_input.txt" - with open(tmp, "w") as f: - f.write( - """ - #[version = "0.0.5"] - def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32], - %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32]) { - %0 = add(%a, %b); - %1 = add(%c, %d); - """ - ) - end = 1555 - for i in range(2, end): - s1 = "\n\t" + "%" + str(i) + " = add(%" + str(i - 1) + ", %" + str(i - 2) + ");" - f.write(s1) - f.write("\n\t" + "add(%" + str(end - 1) + ", %" + str(end - 2) + ")") - f.write("\n\t}") - - with open(tmp) as f: - mod = f.read() - f.close() - - os.remove(tmp) - - return tvm.parser.parse( - mod, - "from_string", - None, - metatable, - ) - config = tvm.target.make_compilation_config(CTXT, TARGETS, HOST_TARGET) - actual_mod = relay.transform.InferType()(in_mod) - actual_mod = relay.transform.PlanDevices(config)(actual_mod) - relay.transform.InferType()(actual_mod) - - if __name__ == "__main__": import sys import pytest From c095f0752d84d9a35c0c6d6f69d5714c3e81360e Mon Sep 17 00:00:00 2001 From: Mei Ye Date: Tue, 8 Feb 2022 00:11:02 +0000 Subject: [PATCH 3/3] [PyTorch] Add grid_sample with zeros and border padding mode for PyTorch. --- include/tvm/relay/attrs/image.h | 6 ++ python/tvm/relay/frontend/pytorch.py | 20 +++++ python/tvm/relay/op/image/_image.py | 3 +- python/tvm/relay/op/image/image.py | 10 ++- python/tvm/topi/image/grid_sample.py | 24 ++++-- python/tvm/topi/testing/grid_sample_python.py | 73 ++++++++++++++----- src/relay/op/image/grid_sample.cc | 3 +- tests/python/frontend/pytorch/test_forward.py | 45 +++++++++++- tests/python/relay/test_op_level5.py | 12 ++- tests/python/topi/python/test_topi_image.py | 11 ++- 10 files changed, 165 insertions(+), 42 deletions(-) diff --git a/include/tvm/relay/attrs/image.h b/include/tvm/relay/attrs/image.h index 78687b38ee76..be207a2d0593 100644 --- a/include/tvm/relay/attrs/image.h +++ b/include/tvm/relay/attrs/image.h @@ -275,6 +275,7 @@ struct AffineGridAttrs : public tvm::AttrsNode { struct GridSampleAttrs : public tvm::AttrsNode { String method; String layout; + String padding_mode; TVM_DECLARE_ATTRS(GridSampleAttrs, "relay.attrs.GridSampleAttrs") { TVM_ATTR_FIELD(method) @@ -287,6 +288,11 @@ struct GridSampleAttrs : public tvm::AttrsNode { "'N', 'C', 'H', 'W' stands for batch, channel, height, and width" "dimensions respectively. Resize is applied on the 'H' and" "'W' dimensions."); + TVM_ATTR_FIELD(padding_mode) + .set_default("zeros") + .describe( + "Specify the padding mode to use." + "zeros, border etc."); } }; diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 217645d817e0..fe3e77afd271 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2896,6 +2896,25 @@ def mv(self, inputs, _): # Chop off the extra result dimension return _op.transform.squeeze(dense_result) + def grid_sampler(self, inputs, input_types): + if inputs[2] == 0: + mode = "bilinear" + else: + msg = "Only bilinear mode is supported in grid_sampler" + raise NotImplementedError(msg) + + if inputs[3] == 0: + padding_mode = "zeros" + elif inputs[3] == 1: + padding_mode = "border" + else: + msg = "Only zeros and border padding mode are supported in grid_sampler" + raise NotImplementedError(msg) + + axes = [0, 3, 1, 2] + grid = _op.transform.transpose(inputs[1], axes) + return _op.image.grid_sample(inputs[0], grid, mode, "NCHW", padding_mode) + # Operator mappings def create_convert_map(self): self.convert_map = { @@ -3124,6 +3143,7 @@ def create_convert_map(self): "aten::einsum": self.einsum, "aten::dot": self.dot, "aten::mv": self.mv, + "aten::grid_sampler": self.grid_sampler, } def update_convert_map(self, custom_map): diff --git a/python/tvm/relay/op/image/_image.py b/python/tvm/relay/op/image/_image.py index d9929869beb1..ec25198adf68 100644 --- a/python/tvm/relay/op/image/_image.py +++ b/python/tvm/relay/op/image/_image.py @@ -365,7 +365,8 @@ def affine_grid_func(attrs, inputs, _): def compute_grid_sample(attrs, inputs, out_dtype): method = attrs.method layout = attrs.layout - return [topi.image.grid_sample(inputs[0], inputs[1], method, layout)] + padding_mode = attrs.padding_mode + return [topi.image.grid_sample(inputs[0], inputs[1], method, layout, padding_mode)] reg.register_injective_schedule("image.grid_sample") diff --git a/python/tvm/relay/op/image/image.py b/python/tvm/relay/op/image/image.py index 30f8a9c21118..eb6c316402c6 100644 --- a/python/tvm/relay/op/image/image.py +++ b/python/tvm/relay/op/image/image.py @@ -455,7 +455,7 @@ def affine_grid(data, target_shape=None): return _make.affine_grid(data, target_shape) -def grid_sample(data, grid, method="bilinear", layout="NCHW"): +def grid_sample(data, grid, method="bilinear", layout="NCHW", padding_mode="zeros"): """Applies bilinear sampling to input feature map. Given :math:`data` and :math:`grid`, then the output is computed by @@ -468,7 +468,8 @@ def grid_sample(data, grid, method="bilinear", layout="NCHW"): :math:`x_{dst}`, :math:`y_{dst}` enumerate all spatial locations in :math:`output`, and :math:`G()` denotes the interpolation function. - The out-boundary points will be padded with zeros. The shape of the output will be + The out-boundary points will be padded with zeros if padding_mode is "zeros". + The shape of the output will be (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]). The operator assumes that :math:`grid` has been normalized to [-1, 1]. @@ -489,9 +490,12 @@ def grid_sample(data, grid, method="bilinear", layout="NCHW"): layout : str The layout of input data and the output. + padding_mode : str + The padding mode for outside grid values. + Returns ------- Output : tvm.Tensor 4-D with shape [batch, 2, out_height, out_width] """ - return _make.grid_sample(data, grid, method, layout) + return _make.grid_sample(data, grid, method, layout, padding_mode) diff --git a/python/tvm/topi/image/grid_sample.py b/python/tvm/topi/image/grid_sample.py index 19a69ef83efe..e3a6dd80405a 100644 --- a/python/tvm/topi/image/grid_sample.py +++ b/python/tvm/topi/image/grid_sample.py @@ -59,7 +59,7 @@ def _compute(n, dim, i, j): return te.compute(oshape, _compute, tag="affine_grid") -def grid_sample(data, grid, method="bilinear", layout="NCHW"): +def grid_sample(data, grid, method="bilinear", layout="NCHW", padding_mode="zeros"): """Applies bilinear sampling to input feature map. Given :math:`data` and :math:`grid`, assuming NCHW layout, then the output is computed by @@ -72,7 +72,8 @@ def grid_sample(data, grid, method="bilinear", layout="NCHW"): :math:`x_{dst}`, :math:`y_{dst}` enumerate all spatial locations in :math:`output`, and :math:`G()` denotes the interpolation method. - The out-boundary points will be padded with zeros. The shape of the output will be + The out-boundary points will be padded with zeros if the padding_mode is "zeros". + The shape of the output will be (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]). The operator assumes that :math:`grid` has been normalized to [-1, 1]. @@ -96,7 +97,7 @@ def grid_sample(data, grid, method="bilinear", layout="NCHW"): Returns ------- Output : tvm.Tensor - 4-D with shape [batch, 2, out_height, out_width] + 4-D with shape [batch, in_channel, out_height, out_width] """ batch, in_channel, in_height, in_width = data.shape out_height, out_width = grid.shape[2:] @@ -104,11 +105,18 @@ def grid_sample(data, grid, method="bilinear", layout="NCHW"): assert layout == "NCHW", "Only NCHW is supported" def _get_pixel_value(n, c, h, w): - return te.if_then_else( - te.all(h >= 0, w >= 0, h < in_height, w < in_width), - data[n, c, h, w], - tir.const(0.0, dtype=data.dtype), - ) + if padding_mode == "zeros": + return te.if_then_else( + te.all(h >= 0, w >= 0, h < in_height, w < in_width), + data[n, c, h, w], + tir.const(0.0, dtype=data.dtype), + ) + if padding_mode == "border": + h_b = te.max(te.min(h, in_height - 1), 0) + w_b = te.max(te.min(w, in_width - 1), 0) + return data[n, c, h_b, w_b] + + raise AssertionError("unsupported padding_mode") def _bilinear_sample(n, c, h, w): x = grid[n, 0, h, w] diff --git a/python/tvm/topi/testing/grid_sample_python.py b/python/tvm/topi/testing/grid_sample_python.py index a8c304f30160..e6b0bef38685 100644 --- a/python/tvm/topi/testing/grid_sample_python.py +++ b/python/tvm/topi/testing/grid_sample_python.py @@ -29,7 +29,7 @@ def affine_grid_python(data, target_shape): return data.reshape(-1, 3).dot(grid).reshape(data.shape[0], 2, *target_shape) -def _bilinear_sample_nchw_python(data, grid): +def _bilinear_sample_nchw_python(data, grid, padding_mode): batch, in_channel, in_height, in_width = data.shape _, _, out_height, out_width = grid.shape out = np.zeros((batch, in_channel, out_height, out_width), dtype=data.dtype) @@ -37,28 +37,63 @@ def _bilinear_sample_nchw_python(data, grid): def _within_bound(y, x): return 0 <= y < in_height and 0 <= x < in_width - for n in range(0, batch): - for h in range(0, out_height): - for w in range(0, out_width): - x, y = grid[n, :, h, w] - y = (y + 1) * (in_height - 1) / 2 - x = (x + 1) * (in_width - 1) / 2 - y0 = int(math.floor(y)) - x0 = int(math.floor(x)) - y1 = y0 + 1 - x1 = x0 + 1 - if _within_bound(y0, x0): - out[n, :, h, w] += data[n, :, y0, x0] * (1.0 - (y - y0)) * (1.0 - (x - x0)) - if _within_bound(y0, x1): + def compute_padding_mode_zeros(): + for n in range(0, batch): + for h in range(0, out_height): + for w in range(0, out_width): + x, y = grid[n, :, h, w] + y = (y + 1) * (in_height - 1) / 2 + x = (x + 1) * (in_width - 1) / 2 + y0 = int(math.floor(y)) + x0 = int(math.floor(x)) + y1 = y0 + 1 + x1 = x0 + 1 + if _within_bound(y0, x0): + out[n, :, h, w] += data[n, :, y0, x0] * (1.0 - (y - y0)) * (1.0 - (x - x0)) + if _within_bound(y0, x1): + out[n, :, h, w] += data[n, :, y0, x1] * (1.0 - (y - y0)) * (x - x0) + if _within_bound(y1, x0): + out[n, :, h, w] += data[n, :, y1, x0] * (y - y0) * (1.0 - (x - x0)) + if _within_bound(y1, x1): + out[n, :, h, w] += data[n, :, y1, x1] * (y - y0) * (x - x0) + + return out + + def get_pixel_value(x, x_max): + return max(min(x, x_max - 1), 0) + + def compute_padding_mode_border(): + for n in range(0, batch): + for h in range(0, out_height): + for w in range(0, out_width): + x, y = grid[n, :, h, w] + y = (y + 1) * (in_height - 1) / 2 + x = (x + 1) * (in_width - 1) / 2 + y0 = int(math.floor(y)) + x0 = int(math.floor(x)) + y1 = y0 + 1 + x1 = x0 + 1 + y0 = get_pixel_value(y0, in_height) + y1 = get_pixel_value(y1, in_height) + x0 = get_pixel_value(x0, in_width) + x1 = get_pixel_value(x1, in_width) + out[n, :, h, w] = data[n, :, y0, x0] * (1.0 - (y - y0)) * (1.0 - (x - x0)) out[n, :, h, w] += data[n, :, y0, x1] * (1.0 - (y - y0)) * (x - x0) - if _within_bound(y1, x0): out[n, :, h, w] += data[n, :, y1, x0] * (y - y0) * (1.0 - (x - x0)) - if _within_bound(y1, x1): out[n, :, h, w] += data[n, :, y1, x1] * (y - y0) * (x - x0) - return out + return out + + if padding_mode == "zeros": + return compute_padding_mode_zeros() + if padding_mode == "border": + return compute_padding_mode_border() -def grid_sample_nchw_python(data, grid, method="bilinear"): + raise ValueError("invalid padding_mode") + + +def grid_sample_nchw_python(data, grid, method="bilinear", padding_mode="zeros"): if method == "bilinear": - return _bilinear_sample_nchw_python(data, grid) + return _bilinear_sample_nchw_python(data, grid, padding_mode) + raise ValueError("invalid method") diff --git a/src/relay/op/image/grid_sample.cc b/src/relay/op/image/grid_sample.cc index 38d4c8103ff2..e0282cc2e8c7 100644 --- a/src/relay/op/image/grid_sample.cc +++ b/src/relay/op/image/grid_sample.cc @@ -116,10 +116,11 @@ bool GridSampleRel(const Array& types, int num_inputs, const Attrs& attrs, // Positional relay function to create affine_grid operator // used by frontend FFI. -Expr MakeGridSample(Expr data, Expr grid, String method, String layout) { +Expr MakeGridSample(Expr data, Expr grid, String method, String layout, String padding_mode) { auto attrs = make_object(); attrs->method = std::move(method); attrs->layout = std::move(layout); + attrs->padding_mode = std::move(padding_mode); static const Op& op = Op::Get("image.grid_sample"); return Call(op, {data, grid}, Attrs(attrs), {}); } diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index c8077f616835..97ef08f7b8a9 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -216,9 +216,8 @@ def verify_model( if not tvm.runtime.enabled(target): continue dev = tvm.device(target, 0) - relay_graph, relay_lib, relay_params = relay.build(mod, target=target, params=params) - relay_model = graph_executor.create(relay_graph, relay_lib, dev) - relay_model.set_input(**relay_params) + lib = relay.build(mod, target=target, params=params) + relay_model = graph_executor.GraphModule(lib["default"](dev)) for name, inp in compiled_input.items(): relay_model.set_input(name, inp) relay_model.run() @@ -4079,5 +4078,45 @@ def test_fn(m, v): verify_model(test_fn, [torch.randn(3, 8), torch.randn(8)]) +def test_grid_sample(): + class Grid_sample_zeros(Module): + def forward(self, x, y): + return torch.nn.functional.grid_sample( + input=x, grid=y, mode="bilinear", padding_mode="zeros", align_corners=True + ) + + class Grid_sample_border(Module): + def forward(self, x, y): + return torch.nn.functional.grid_sample( + input=x, grid=y, mode="bilinear", padding_mode="border", align_corners=True + ) + + data = torch.rand([4, 4, 16, 32]).float() + grid = torch.rand([4, 8, 8, 2]).float() + verify_model(Grid_sample_zeros(), input_data=[data, grid]) + verify_model(Grid_sample_border(), input_data=[data, grid]) + + +def test_list_tuple(): + """test compilation error for a Python list followed by a prim::TupleConstruct.""" + + class List_tuple(Module): + def forward(self, x): + merged = [] + mask_list = [] + for i in range(3): + w0 = torch.sigmoid(x) + merged.append((w0, w0)) + mask_list.append(x) + + for i in range(3): + merged[i] = merged[i][0] + merged[i][1] + return mask_list[2], merged + + x = torch.rand([4, 4, 16, 32]).float() + script_module = torch.jit.trace(List_tuple(), x, strict=False).eval() + mod, params = relay.frontend.from_pytorch(script_module, [("x", x.shape)]) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index 6ad3b3827dce..2f1b020c5450 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -1392,20 +1392,24 @@ def verify_affine_grid(num_batch, target_shape): @tvm.testing.uses_gpu def test_grid_sample(): - def verify_grid_sample(data_shape, grid_shape): + def verify_grid_sample(data_shape, grid_shape, padding_mode="zeros"): dtype = "float32" batch, channel, _, _ = data_shape _, _, out_height, out_width = grid_shape data = relay.var("data", relay.ty.TensorType(data_shape, dtype)) grid = relay.var("grid", relay.ty.TensorType(grid_shape, dtype)) - y = relay.image.grid_sample(data, grid, method="bilinear", layout="NCHW") + y = relay.image.grid_sample( + data, grid, method="bilinear", layout="NCHW", padding_mode=padding_mode + ) yy = run_infer_type(y) assert yy.checked_type == relay.TensorType((batch, channel, out_height, out_width), dtype) func = relay.Function([data, grid], y) data_np = np.random.uniform(size=data_shape).astype(dtype) grid_np = np.random.uniform(size=grid_shape, low=-1.5, high=1.5).astype(dtype) - ref_res = tvm.topi.testing.grid_sample_nchw_python(data_np, grid_np, method="bilinear") + ref_res = tvm.topi.testing.grid_sample_nchw_python( + data_np, grid_np, method="bilinear", padding_mode=padding_mode + ) for target, dev in tvm.testing.enabled_targets(): for kind in ["graph", "debug"]: @@ -1416,6 +1420,8 @@ def verify_grid_sample(data_shape, grid_shape): verify_grid_sample((4, 4, 16, 32), (4, 2, 8, 8)) verify_grid_sample((4, 4, 16, 32), (4, 2, 32, 32)) + verify_grid_sample((4, 4, 16, 32), (4, 2, 8, 8), "border") + verify_grid_sample((4, 4, 16, 32), (4, 2, 32, 32), "border") @tvm.testing.uses_gpu diff --git a/tests/python/topi/python/test_topi_image.py b/tests/python/topi/python/test_topi_image.py index 062cf7964ecb..9f4b67354075 100644 --- a/tests/python/topi/python/test_topi_image.py +++ b/tests/python/topi/python/test_topi_image.py @@ -274,18 +274,20 @@ def check_target(target, dev): @tvm.testing.uses_gpu def test_grid_sample(): - def verify_grid_sample(data_shape, grid_shape): + def verify_grid_sample(data_shape, grid_shape, padding_mode="zeros"): dtype = "float32" data = te.placeholder(data_shape, dtype=dtype) grid = te.placeholder(grid_shape, dtype=dtype) - out = topi.image.grid_sample(data, grid, "bilinear") + out = topi.image.grid_sample(data, grid, "bilinear", padding_mode=padding_mode) @memoize("topi.tests.test_grid_sample.verify_grid_sample") def get_ref_data(): data_np = np.random.uniform(size=data_shape).astype(dtype) # allow grid values to be out-of-bound grid_np = np.random.uniform(size=grid_shape, low=-1.5, high=1.5).astype(dtype) - out_np = tvm.topi.testing.grid_sample_nchw_python(data_np, grid_np, "bilinear") + out_np = tvm.topi.testing.grid_sample_nchw_python( + data_np, grid_np, "bilinear", padding_mode + ) return data_np, grid_np, out_np data_np, grid_np, out_np = get_ref_data() @@ -306,7 +308,8 @@ def check_target(target, dev): check_target(target, dev) verify_grid_sample((4, 4, 16, 32), (4, 2, 8, 8)) - verify_grid_sample((4, 4, 16, 32), (4, 2, 32, 32)) + verify_grid_sample((4, 4, 16, 32), (4, 2, 32, 32), "border") + verify_grid_sample((4, 4, 16, 32), (4, 2, 8, 8), "border") if __name__ == "__main__":