diff --git a/backends/cortex_m/test/ops/test_conv.py b/backends/cortex_m/test/ops/test_conv.py index f2ac0df0cc2..c3b7863dc7b 100644 --- a/backends/cortex_m/test/ops/test_conv.py +++ b/backends/cortex_m/test/ops/test_conv.py @@ -83,7 +83,6 @@ class CortexMConv3D(torch.nn.Module): def __init__(self, *args, **kwargs): super().__init__() self.conv = torch.nn.Conv3d(*args, **kwargs, bias=False) - self.conv.weight.data.fill_(2.0) def forward(self, x): return self.conv(x) @@ -223,13 +222,12 @@ def forward(self, x): model=CortexMConv1D(1, 1, 1), example_inputs=(ramp_tensor(0, 10, (1, 3, 2)),), ), + # Contiguous on purpose: the Arm test runner only reorders NHWC inputs and its + # own 5-D order, so a channels_last_3d input would reach the simulator in the + # wrong byte order. "conv3d": McuTestCase( - model=CortexMConv3D(1, 1, 1), - example_inputs=( - ramp_tensor(-1000, 1000, (2, 1, 3, 3, 3)).to( - memory_format=torch.channels_last_3d - ), - ), + model=CortexMConv3D(2, 4, (2, 3, 3), stride=(1, 2, 2), padding=1), + example_inputs=(ramp_tensor(-10, 10, (2, 2, 4, 6, 6)),), ), "conv2d_x3": McuTestCase( model=CortexMConv2Dx3(), @@ -349,7 +347,6 @@ def test_dialect_conv2d(test_case, cortex_m_target): xfails_implementation: dict[str, xfail_type] = { "conv1d": "Currently not supported.", - "conv3d": "Currently not supported.", } diff --git a/backends/test/suite/runner.py b/backends/test/suite/runner.py index e90e951db20..f509a75b3de 100644 --- a/backends/test/suite/runner.py +++ b/backends/test/suite/runner.py @@ -37,7 +37,6 @@ TestResult, ) from executorch.exir import EdgeProgramManager -from executorch.exir.dialects._ops import ops as exir_ops # A list of all runnable test suites and the corresponding python package. @@ -47,24 +46,6 @@ } -def _graph_has_unsupported_patterns(program: torch.export.ExportedProgram) -> bool: - # Returns true if the model contains patterns that will fail when running on the ET - # portable kernel library. - - # Check for 3d convolutions. All convs (1d, 2d, 3d) use the same op, so we need to look at - # the input meta to determine the rank. - for node in program.graph.nodes: - if ( - node.op == "call_function" - and node.target == exir_ops.edge.aten.convolution.default - ): - in_rank = node.args[0].meta["val"].dim() - if in_rank > 4: - return True - - return False - - def _get_test_seed(test_base_name: str) -> int: # Set the seed based on the test base name to give consistent inputs between backends. Add the # run seed to allow for reproducible results, but still allow for run-to-run variation. @@ -185,7 +166,7 @@ def build_result( # Check if any undelegated ops are in the unsupported ops set. has_unsupported_ops = any( op in UNSUPPORTED_PORTABLE_OPS for op in undelegated_op_counts.keys() - ) or _graph_has_unsupported_patterns(edge_manager._etrecord.edge_dialect_program) + ) # Skip the test if there are unsupported portable ops remaining. if has_unsupported_ops: diff --git a/kernels/portable/cpu/op_convolution.cpp b/kernels/portable/cpu/op_convolution.cpp index 25186b92121..3a8a2f40014 100644 --- a/kernels/portable/cpu/op_convolution.cpp +++ b/kernels/portable/cpu/op_convolution.cpp @@ -12,7 +12,6 @@ #include #include -#include #include #include @@ -24,8 +23,6 @@ using Tensor = executorch::aten::Tensor; using ScalarType = executorch::aten::ScalarType; using IntArrayRef = executorch::aten::ArrayRef; using SizesArrayRef = executorch::aten::ArrayRef; -using DimOrderArrayRef = - executorch::aten::ArrayRef; using StridesArrayRef = executorch::aten::ArrayRef; @@ -36,7 +33,9 @@ struct IndexRange { int64_t end; }; -IndexRange transposed_kernel_range( +// The range helpers do not depend on the dtype, so a single out-of-line copy +// serves every dtype and spatial dim. +ET_NOINLINE IndexRange transposed_kernel_range( int64_t begin, int64_t end, int64_t input_size, @@ -52,7 +51,7 @@ IndexRange transposed_kernel_range( std::min(kernel_size, (end - 1 + padding) / dilation + 1)}; } -IndexRange transposed_input_range( +ET_NOINLINE IndexRange transposed_input_range( int64_t begin, int64_t end, int64_t input_size, @@ -65,7 +64,7 @@ IndexRange transposed_input_range( // Keep the tile buffer out of the regular convolution's stack frame. template -ET_NOINLINE void transposed_conv2d_impl( +ET_NOINLINE void transposed_conv3d_impl( const CTYPE* const in_ptr, SizesArrayRef in_sizes, StridesArrayRef in_strides, @@ -88,13 +87,17 @@ ET_NOINLINE void transposed_conv2d_impl( size_t in_C = in_sizes[1]; size_t out_C = out_sizes[1]; - int64_t out_H = out_sizes[2]; - int64_t in_H = in_sizes[2]; - int64_t w_H = w_sizes[2]; + int64_t out_D = out_sizes[2]; + int64_t in_D = in_sizes[2]; + int64_t w_D = w_sizes[2]; - int64_t out_W = out_sizes[3]; - int64_t in_W = in_sizes[3]; - int64_t w_W = w_sizes[3]; + int64_t out_H = out_sizes[3]; + int64_t in_H = in_sizes[3]; + int64_t w_H = w_sizes[3]; + + int64_t out_W = out_sizes[4]; + int64_t in_W = in_sizes[4]; + int64_t w_W = w_sizes[4]; size_t in_C_per_group = in_C / groups; size_t in_c_start = group * in_C_per_group; @@ -102,73 +105,94 @@ ET_NOINLINE void transposed_conv2d_impl( size_t out_C_per_group = out_C / groups; size_t out_c_start = group * out_C_per_group; - const int64_t stride_y = val_at(stride, 0); - const int64_t padding_y = val_at(padding, 0, /*default_value=*/0); - const int64_t dilation_y = val_at(dilation, 0); - const int64_t stride_x = val_at(stride, 1); - const int64_t padding_x = val_at(padding, 1, /*default_value=*/0); - const int64_t dilation_x = val_at(dilation, 1); + const int64_t stride_z = stride[0]; + const int64_t padding_z = padding[0]; + const int64_t dilation_z = dilation[0]; + const int64_t stride_y = stride[1]; + const int64_t padding_y = padding[1]; + const int64_t dilation_y = dilation[1]; + const int64_t stride_x = stride[2]; + const int64_t padding_x = padding[2]; + const int64_t dilation_x = dilation[2]; using COMPUTE_T = typename executorch::runtime::promote_types::type; // Reuse weights across a bounded tile, narrowing only after reduction. constexpr int64_t kTileHeight = 8; constexpr int64_t kTileWidth = 32; - for (int64_t tile_y = 0; tile_y < out_H; tile_y += kTileHeight) { - const int64_t tile_y_end = std::min(out_H, tile_y + kTileHeight); - const auto kernel_y = transposed_kernel_range( - tile_y, tile_y_end, in_H, w_H, stride_y, padding_y, dilation_y); - for (int64_t tile_x = 0; tile_x < out_W; tile_x += kTileWidth) { - const int64_t tile_x_end = std::min(out_W, tile_x + kTileWidth); - const auto kernel_x = transposed_kernel_range( - tile_x, tile_x_end, in_W, w_W, stride_x, padding_x, dilation_x); - COMPUTE_T accum[kTileHeight * kTileWidth] = {}; - for (int64_t w_y = kernel_y.begin; w_y < kernel_y.end; ++w_y) { - const int64_t offset_y = w_y * dilation_y - padding_y; - const auto input_y = transposed_input_range( - tile_y, tile_y_end, in_H, offset_y, stride_y); - if (input_y.begin >= input_y.end) { - continue; - } - for (int64_t w_x = kernel_x.begin; w_x < kernel_x.end; ++w_x) { - const int64_t offset_x = w_x * dilation_x - padding_x; - const auto input_x = transposed_input_range( - tile_x, tile_x_end, in_W, offset_x, stride_x); - if (input_x.begin >= input_x.end) { + for (int64_t out_z = 0; out_z < out_D; ++out_z) { + const auto kernel_z = transposed_kernel_range( + out_z, out_z + 1, in_D, w_D, stride_z, padding_z, dilation_z); + for (int64_t tile_y = 0; tile_y < out_H; tile_y += kTileHeight) { + const int64_t tile_y_end = std::min(out_H, tile_y + kTileHeight); + const auto kernel_y = transposed_kernel_range( + tile_y, tile_y_end, in_H, w_H, stride_y, padding_y, dilation_y); + for (int64_t tile_x = 0; tile_x < out_W; tile_x += kTileWidth) { + const int64_t tile_x_end = std::min(out_W, tile_x + kTileWidth); + const auto kernel_x = transposed_kernel_range( + tile_x, tile_x_end, in_W, w_W, stride_x, padding_x, dilation_x); + COMPUTE_T accum[kTileHeight * kTileWidth] = {}; + for (int64_t w_z = kernel_z.begin; w_z < kernel_z.end; ++w_z) { + // A tile is one output deep, so at most one input depth lands on it. + const auto input_z = transposed_input_range( + out_z, out_z + 1, in_D, w_z * dilation_z - padding_z, stride_z); + if (input_z.begin >= input_z.end) { continue; } - for (const auto in_c : - c10::irange(in_c_start, in_c_start + in_C_per_group)) { - const size_t w_idx = in_c * w_strides[0] + - (out_c - out_c_start) * w_strides[1] + w_y * w_strides[2] + - w_x * w_strides[3]; - const COMPUTE_T w_val = w_ptr[w_idx]; - int64_t out_y = input_y.begin * stride_y + offset_y - tile_y; - for (int64_t in_y = input_y.begin; in_y < input_y.end; ++in_y) { - size_t in_idx = batch * in_strides[0] + in_c * in_strides[1] + - in_y * in_strides[2] + input_x.begin * in_strides[3]; - int64_t out_x = input_x.begin * stride_x + offset_x - tile_x; - for (int64_t in_x = input_x.begin; in_x < input_x.end; ++in_x) { - const COMPUTE_T in_val = in_ptr[in_idx]; - accum[out_y * kTileWidth + out_x] += in_val * w_val; - in_idx += in_strides[3]; - out_x += stride_x; + const size_t in_z_idx = + batch * in_strides[0] + input_z.begin * in_strides[2]; + const size_t w_z_idx = + (out_c - out_c_start) * w_strides[1] + w_z * w_strides[2]; + for (int64_t w_y = kernel_y.begin; w_y < kernel_y.end; ++w_y) { + const int64_t offset_y = w_y * dilation_y - padding_y; + const auto input_y = transposed_input_range( + tile_y, tile_y_end, in_H, offset_y, stride_y); + if (input_y.begin >= input_y.end) { + continue; + } + for (int64_t w_x = kernel_x.begin; w_x < kernel_x.end; ++w_x) { + const int64_t offset_x = w_x * dilation_x - padding_x; + const auto input_x = transposed_input_range( + tile_x, tile_x_end, in_W, offset_x, stride_x); + if (input_x.begin >= input_x.end) { + continue; + } + for (const auto in_c : + c10::irange(in_c_start, in_c_start + in_C_per_group)) { + const size_t w_idx = w_z_idx + in_c * w_strides[0] + + w_y * w_strides[3] + w_x * w_strides[4]; + const COMPUTE_T w_val = w_ptr[w_idx]; + int64_t out_y = input_y.begin * stride_y + offset_y - tile_y; + for (int64_t in_y = input_y.begin; in_y < input_y.end; ++in_y) { + size_t in_idx = in_z_idx + in_c * in_strides[1] + + in_y * in_strides[3] + input_x.begin * in_strides[4]; + int64_t out_x = input_x.begin * stride_x + offset_x - tile_x; + for (int64_t in_x = input_x.begin; in_x < input_x.end; + ++in_x) { + const COMPUTE_T in_val = in_ptr[in_idx]; + accum[out_y * kTileWidth + out_x] += in_val * w_val; + in_idx += in_strides[4]; + out_x += stride_x; + } + out_y += stride_y; + } } - out_y += stride_y; } } } - } - const COMPUTE_T bias_val = bias_ptr == nullptr - ? COMPUTE_T{0} - : load_bias(&bias_ptr[out_c * bias.value().element_size()]); - for (int64_t out_y = tile_y; out_y < tile_y_end; ++out_y) { - size_t out_idx = batch * out_strides[0] + out_c * out_strides[1] + - out_y * out_strides[2] + tile_x * out_strides[3]; - for (int64_t out_x = tile_x; out_x < tile_x_end; ++out_x) { - out_ptr[out_idx] = - accum[(out_y - tile_y) * kTileWidth + out_x - tile_x] + bias_val; - out_idx += out_strides[3]; + const COMPUTE_T bias_val = bias_ptr == nullptr + ? COMPUTE_T{0} + : load_bias(&bias_ptr[out_c * bias.value().element_size()]); + for (int64_t out_y = tile_y; out_y < tile_y_end; ++out_y) { + size_t out_idx = batch * out_strides[0] + out_c * out_strides[1] + + out_z * out_strides[2] + out_y * out_strides[3] + + tile_x * out_strides[4]; + for (int64_t out_x = tile_x; out_x < tile_x_end; ++out_x) { + out_ptr[out_idx] = + accum[(out_y - tile_y) * kTileWidth + out_x - tile_x] + + bias_val; + out_idx += out_strides[4]; + } } } } @@ -176,7 +200,7 @@ ET_NOINLINE void transposed_conv2d_impl( } template -void conv2d_impl( +void conv3d_impl( const CTYPE* const in_ptr, SizesArrayRef in_sizes, StridesArrayRef in_strides, @@ -198,7 +222,7 @@ void conv2d_impl( const size_t out_c, bool transposed) { if (transposed) { - transposed_conv2d_impl( + transposed_conv3d_impl( in_ptr, in_sizes, in_strides, @@ -222,78 +246,102 @@ void conv2d_impl( } size_t in_C = in_sizes[1]; - int64_t out_H = out_sizes[2]; - int64_t in_H = in_sizes[2]; - int64_t w_H = w_sizes[2]; + int64_t out_D = out_sizes[2]; + int64_t in_D = in_sizes[2]; + int64_t w_D = w_sizes[2]; - int64_t out_W = out_sizes[3]; - int64_t in_W = in_sizes[3]; - int64_t w_W = w_sizes[3]; + int64_t out_H = out_sizes[3]; + int64_t in_H = in_sizes[3]; + int64_t w_H = w_sizes[3]; + + int64_t out_W = out_sizes[4]; + int64_t in_W = in_sizes[4]; + int64_t w_W = w_sizes[4]; size_t in_C_per_group = in_C / groups; size_t in_c_start = group * in_C_per_group; - executorch::aten::SizesType in_coord[kTensorDimensionLimit]; - in_coord[0] = batch; - executorch::aten::SizesType out_coord[kTensorDimensionLimit]; - out_coord[0] = batch; - out_coord[1] = out_c; - executorch::aten::SizesType w_coord[kTensorDimensionLimit]; - - const int64_t stride_y = val_at(stride, 0); - const int64_t padding_y = val_at(padding, 0, /*default_value=*/0); - const int64_t dilation_y = val_at(dilation, 0); - const int64_t stride_x = val_at(stride, 1); - const int64_t padding_x = val_at(padding, 1, /*default_value=*/0); - const int64_t dilation_x = val_at(dilation, 1); + const int64_t stride_z = stride[0]; + const int64_t padding_z = padding[0]; + const int64_t dilation_z = dilation[0]; + const int64_t stride_y = stride[1]; + const int64_t padding_y = padding[1]; + const int64_t dilation_y = dilation[1]; + const int64_t stride_x = stride[2]; + const int64_t padding_x = padding[2]; + const int64_t dilation_x = dilation[2]; using COMPUTE_T = typename executorch::runtime::promote_types::type; - - for (const auto out_y : c10::irange(out_H)) { - out_coord[2] = out_y; - for (const auto out_x : c10::irange(out_W)) { - out_coord[3] = out_x; - COMPUTE_T accum = 0; - for (const auto in_c : - c10::irange(in_c_start, in_c_start + in_C_per_group)) { - in_coord[1] = in_c; - w_coord[0] = out_c; - w_coord[1] = in_c - in_c_start; - for (const auto w_y : c10::irange(w_H)) { - w_coord[2] = w_y; - int64_t in_y = stride_y * out_y + dilation_y * w_y - padding_y; - if (in_y < 0 || in_y >= in_H) { + const COMPUTE_T bias_val = bias_ptr == nullptr + ? COMPUTE_T{0} + : load_bias(&bias_ptr[out_c * bias.value().element_size()]); + + for (const auto out_z : c10::irange(out_D)) { + for (const auto out_y : c10::irange(out_H)) { + for (const auto out_x : c10::irange(out_W)) { + COMPUTE_T accum = 0; + // Depth taps sit outside the channel loop, so for 1-D and 2-D inputs + // the depth loop runs once and adds no work per channel. + for (const auto w_z : c10::irange(w_D)) { + int64_t in_z = stride_z * out_z + dilation_z * w_z - padding_z; + if (in_z < 0 || in_z >= in_D) { continue; } - in_coord[2] = in_y; - for (const auto w_x : c10::irange(w_W)) { - w_coord[3] = w_x; - int64_t in_x = stride_x * out_x + dilation_x * w_x - padding_x; - if (in_x < 0 || in_x >= in_W) { - continue; + for (const auto in_c : + c10::irange(in_c_start, in_c_start + in_C_per_group)) { + const int64_t in_plane_idx = batch * in_strides[0] + + in_c * in_strides[1] + in_z * in_strides[2]; + const int64_t w_plane_idx = out_c * w_strides[0] + + (in_c - in_c_start) * w_strides[1] + w_z * w_strides[2]; + for (const auto w_y : c10::irange(w_H)) { + int64_t in_y = stride_y * out_y + dilation_y * w_y - padding_y; + if (in_y < 0 || in_y >= in_H) { + continue; + } + int64_t in_x = stride_x * out_x - padding_x; + int64_t in_idx = + in_plane_idx + in_y * in_strides[3] + in_x * in_strides[4]; + int64_t w_idx = w_plane_idx + w_y * w_strides[3]; + for (int64_t w_x = 0; w_x < w_W; ++w_x) { + if (in_x >= 0 && in_x < in_W) { + const COMPUTE_T in_val = in_ptr[in_idx]; + const COMPUTE_T w_val = w_ptr[w_idx]; + accum += in_val * w_val; + } + in_x += dilation_x; + in_idx += dilation_x * in_strides[4]; + w_idx += w_strides[4]; + } } - in_coord[3] = in_x; - const size_t in_idx = - calculate_linear_index(in_coord, in_strides.data(), 4); - const size_t w_idx = - calculate_linear_index(w_coord, w_strides.data(), 4); - const COMPUTE_T in_val = in_ptr[in_idx]; - const COMPUTE_T w_val = w_ptr[w_idx]; - accum += in_val * w_val; } } + const int64_t out_idx = batch * out_strides[0] + + out_c * out_strides[1] + out_z * out_strides[2] + + out_y * out_strides[3] + out_x * out_strides[4]; + out_ptr[out_idx] = accum + bias_val; } - if (bias_ptr != nullptr) { - accum += load_bias(&bias_ptr[out_c * bias.value().element_size()]); - } - const size_t out_idx = - calculate_linear_index(out_coord, out_strides.data(), 4); - out_ptr[out_idx] = accum; } } } +// Views t as (N, C, D, H, W) by inserting leading spatial dims of size 1. +void get_3d_view( + const Tensor& t, + executorch::aten::SizesType* sizes, + executorch::aten::StridesType* strides) { + for (const auto d : c10::irange(5)) { + sizes[d] = 1; + strides[d] = 0; + } + const ssize_t added_dims = 5 - t.dim(); + for (const auto d : c10::irange(t.dim())) { + const ssize_t view_d = d < 2 ? d : d + added_dims; + sizes[view_d] = t.size(d); + strides[view_d] = t.strides()[d]; + } +} + // Keep dtype-specific loops separate from the scalar-type dispatch. template ET_NOINLINE void convolution_wrapper( @@ -307,87 +355,30 @@ ET_NOINLINE void convolution_wrapper( bool transposed, int64_t groups, Tensor& out) { - SizesArrayRef in_sizes = in.sizes(); - SizesArrayRef weight_sizes = weight.sizes(); - SizesArrayRef out_sizes = out.sizes(); - - DimOrderArrayRef in_dim_order = in.dim_order(); - DimOrderArrayRef weight_dim_order = weight.dim_order(); - DimOrderArrayRef out_dim_order = out.dim_order(); - - IntArrayRef stride_ = stride; - IntArrayRef padding_ = padding; - IntArrayRef dilation_ = dilation; - - // Define arrays for modified sizes, etc. which will potentially be used - executorch::aten::SizesType in_sizes_arr[kTensorDimensionLimit]; - executorch::aten::DimOrderType in_dim_order_arr[kTensorDimensionLimit]; - size_t in_ndim; - executorch::aten::SizesType weight_sizes_arr[kTensorDimensionLimit]; - executorch::aten::DimOrderType weight_dim_order_arr[kTensorDimensionLimit]; - size_t weight_ndim; - executorch::aten::SizesType out_sizes_arr[kTensorDimensionLimit]; - executorch::aten::DimOrderType out_dim_order_arr[kTensorDimensionLimit]; - size_t out_ndim; - - int64_t stride_arr[2]; - int64_t padding_arr[2]; - int64_t dilation_arr[2]; - - // If in has a dim of 3, then a 1D convolution will be performed. A 1D - // convolution is equivalent to a 2D convolution where the height dim of - // all tensors is 1, and stride = 1, padding = 0, and dilation = 1 for - // the height dimension. Therefore the tensor sizes are unsqueezed and - // the stride, padding, and dilation are adjusted so that a 2D - // convolution implementation can be used. - if (in.dim() == 3) { - get_unsqueezed_sizes(in, 2, in_sizes_arr, in_ndim); - in_sizes = {in_sizes_arr, in_ndim}; - get_unsqueezed_dim_order(in, 2, in_dim_order_arr); - in_dim_order = {in_dim_order_arr, in_ndim}; - - get_unsqueezed_sizes(weight, 2, weight_sizes_arr, weight_ndim); - weight_sizes = {weight_sizes_arr, weight_ndim}; - get_unsqueezed_dim_order(weight, 2, weight_dim_order_arr); - weight_dim_order = {weight_dim_order_arr, weight_ndim}; - - get_unsqueezed_sizes(out, 2, out_sizes_arr, out_ndim); - out_sizes = {out_sizes_arr, out_ndim}; - get_unsqueezed_dim_order(out, 2, out_dim_order_arr); - out_dim_order = {out_dim_order_arr, out_ndim}; - - stride_arr[0] = 1; - stride_arr[1] = stride[0]; - stride_ = {stride_arr, 2}; - - padding_arr[0] = 0; - padding_arr[1] = padding[0]; - padding_ = {padding_arr, 2}; - - dilation_arr[0] = 1; - if (dilation.size() > 0) { - dilation_arr[1] = dilation[0]; - } else { - dilation_arr[1] = 1; - } - dilation_ = {dilation_arr, 2}; + // A 1-D or 2-D convolution is a 3-D one whose leading spatial dims have size + // 1, stride 1, padding 0 and dilation 1, so every rank shares the 3-D + // implementation. + executorch::aten::SizesType in_sizes[5]; + executorch::aten::StridesType in_strides[5]; + get_3d_view(in, in_sizes, in_strides); + executorch::aten::SizesType weight_sizes[5]; + executorch::aten::StridesType weight_strides[5]; + get_3d_view(weight, weight_sizes, weight_strides); + executorch::aten::SizesType out_sizes[5]; + executorch::aten::StridesType out_strides[5]; + get_3d_view(out, out_sizes, out_strides); + + int64_t stride_3d[] = {1, 1, 1}; + int64_t padding_3d[] = {0, 0, 0}; + int64_t dilation_3d[] = {1, 1, 1}; + const ssize_t spatial_dims = in.dim() - 2; + for (const auto i : c10::irange(spatial_dims)) { + const ssize_t d = 3 - spatial_dims + i; + stride_3d[d] = val_at(stride, i); + padding_3d[d] = val_at(padding, i, /*default_value=*/0); + dilation_3d[d] = val_at(dilation, i); } - executorch::aten::StridesType in_strides[kTensorDimensionLimit]; - dim_order_to_stride_nocheck( - in_sizes.data(), in_dim_order.data(), in_sizes.size(), in_strides); - - executorch::aten::StridesType weight_strides[kTensorDimensionLimit]; - dim_order_to_stride_nocheck( - weight_sizes.data(), - weight_dim_order.data(), - weight_sizes.size(), - weight_strides); - - executorch::aten::StridesType out_strides[kTensorDimensionLimit]; - dim_order_to_stride_nocheck( - out_sizes.data(), out_dim_order.data(), out_sizes.size(), out_strides); - CTYPE* const out_ptr = out.mutable_data_ptr(); const CTYPE* const in_ptr = in.const_data_ptr(); const CTYPE* const w_ptr = weight.const_data_ptr(); @@ -406,23 +397,23 @@ ET_NOINLINE void convolution_wrapper( // Populate all the out channels in the group for (const auto out_c : c10::irange(out_c_start, out_c_start + out_C_per_group)) { - conv2d_impl( + conv3d_impl( in_ptr, in_sizes, - {in_strides, 4}, + in_strides, w_ptr, weight_sizes, - {weight_strides, 4}, + weight_strides, bias, bias_ptr, load_bias, - stride_, - padding_, - dilation_, + stride_3d, + padding_3d, + dilation_3d, groups, out_ptr, out_sizes, - {out_strides, 4}, + out_strides, batch, group, out_c, @@ -464,9 +455,6 @@ Tensor& convolution_out( InvalidArgument, out); - ET_KERNEL_CHECK( - ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); - size_t output_ndim = 0; executorch::aten::SizesType output_sizes[kTensorDimensionLimit]; get_convolution_out_target_size( diff --git a/kernels/portable/cpu/util/kernel_ops_util.cpp b/kernels/portable/cpu/util/kernel_ops_util.cpp index 1e3bbcdb92d..7d6501e08bd 100644 --- a/kernels/portable/cpu/util/kernel_ops_util.cpp +++ b/kernels/portable/cpu/util/kernel_ops_util.cpp @@ -207,8 +207,8 @@ int64_t _kernel_output_size_helper( bool transposed, int64_t output_padding) { if (transposed) { - return (inputSize - 1) * stride - 2 * pad + dilation * (kernelSize - 1) + - output_padding + 1; + return (static_cast(inputSize) - 1) * stride - 2 * pad + + dilation * (kernelSize - 1) + output_padding + 1; } int64_t numerator = inputSize + 2 * pad - dilation * (kernelSize - 1) - 1 + (ceil_mode ? stride - 1 : 0); @@ -452,42 +452,83 @@ bool check_convolution_args( ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(out)); ET_CHECK_OR_RETURN_FALSE( - in.dim() == 3 || in.dim() == 4, - "Expect input tensor to be 3-D or 4-D, but got, %zu.", + in.dim() >= 3 && in.dim() <= 5, + "Expect input tensor to be 3-D, 4-D or 5-D, but got, %zu.", static_cast(in.dim())); ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(weight, in.dim())); ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(out, in.dim())); - if (bias.has_value()) { - ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(bias.value(), 1)); - ET_CHECK_OR_RETURN_FALSE( - bias.value().size(0) == transposed ? groups * weight.size(1) - : weight.size(0), - "bias length must equal number of output channels, but got %" ET_PRI_TENSOR_SIZE - "; expected %" PRId64, - bias.value().size(0), - transposed ? groups * weight.size(1) : weight.size(0)); - } - - int64_t kernel_size[2]; - size_t kernel_ndim = 2; - if (weight.dim() == 3) { - kernel_size[0] = weight.size(2); - kernel_ndim = 1; - } else { - kernel_size[0] = weight.size(2); - kernel_size[1] = weight.size(3); + int64_t kernel_size[kTensorDimensionLimit]; + const size_t kernel_ndim = weight.dim() - 2; + for (const auto i : c10::irange(kernel_ndim)) { + kernel_size[i] = weight.size(i + 2); } + ET_LOG_AND_RETURN_IF_FALSE( + kernel_size_is_valid({kernel_size, kernel_ndim}, kernel_ndim)); ET_LOG_AND_RETURN_IF_FALSE( stride_is_valid(stride, kernel_ndim, /*allow_empty=*/false)); ET_LOG_AND_RETURN_IF_FALSE( padding_is_valid(padding, {kernel_size, kernel_ndim}, kernel_ndim)); ET_LOG_AND_RETURN_IF_FALSE(dilation_is_valid(dilation, kernel_ndim)); + // Larger values would overflow int64_t in the size math below. + for (const IntArrayRef param : {stride, padding, dilation}) { + for (const int64_t value : param) { + ET_CHECK_OR_RETURN_FALSE( + value <= std::numeric_limits::max(), + "Expected stride, padding and dilation to fit in 32 bits, but got %" PRId64, + value); + } + } if (transposed) { ET_LOG_AND_RETURN_IF_FALSE( output_padding_is_valid(output_padding, stride, dilation, kernel_ndim)); } + for (const auto i : c10::irange(kernel_ndim)) { + const int64_t input_size = in.size(i + 2); + const int64_t pad = val_at(padding, i, /*default_value=*/0); + const int64_t dil = val_at(dilation, i, /*default_value=*/1); + // Like PyTorch, a spatial dim may only be empty when the batch or channel + // dim is empty too. + ET_CHECK_OR_RETURN_FALSE( + input_size > 0 || in.size(0) == 0 || in.size(1) == 0, + "Input spatial dim %zu is empty, but the batch and channel dims are not", + i); + if (!transposed) { + // Like PyTorch, reject a dilated kernel larger than the padded input. The + // output size math can otherwise round such a size up to 1 for strides + // above 1. + const int64_t padded_input_size = input_size + 2 * pad; + const int64_t dilated_kernel_size = dil * (kernel_size[i] - 1) + 1; + ET_CHECK_OR_RETURN_FALSE( + padded_input_size >= dilated_kernel_size, + "Padded input size %" PRId64 + " is smaller than the dilated kernel size %" PRId64 + " in spatial dim %zu", + padded_input_size, + dilated_kernel_size, + i); + } + // calculate_kernel_output_sizes narrows the output sizes to SizesType. + const int64_t output_size = _kernel_output_size_helper( + input_size, + kernel_size[i], + pad, + val_at(stride, i), + dil, + /*ceil_mode=*/false, + transposed, + transposed ? val_at(output_padding, i, /*default_value=*/0) : 0); + ET_CHECK_OR_RETURN_FALSE( + output_size > 0 && + output_size <= + std::numeric_limits::max(), + "Output size %" PRId64 " in spatial dim %zu is out of range", + output_size, + i); + } + ET_CHECK_OR_RETURN_FALSE( + groups > 0, "groups must be positive, but got %" PRId64, groups); ET_CHECK_OR_RETURN_FALSE( weight.size(0) >= groups, "Given groups=%" PRId64 ", expected weight to be at least %" PRId64 @@ -503,6 +544,20 @@ bool check_convolution_args( groups, weight.size(0)); + // The groups checks above keep groups * weight.size(1) from overflowing. + if (bias.has_value()) { + ET_LOG_AND_RETURN_IF_FALSE(tensor_is_rank(bias.value(), 1)); + ET_LOG_AND_RETURN_IF_FALSE( + executorch::runtime::tensor_is_realhbf16_type(bias.value())); + ET_CHECK_OR_RETURN_FALSE( + bias.value().size(0) == + (transposed ? groups * weight.size(1) : weight.size(0)), + "bias length must equal number of output channels, but got %" ET_PRI_TENSOR_SIZE + "; expected %" PRId64, + bias.value().size(0), + transposed ? groups * weight.size(1) : weight.size(0)); + } + if (!transposed) { ET_CHECK_OR_RETURN_FALSE( in.size(1) == groups * weight.size(1), @@ -547,14 +602,10 @@ void get_convolution_out_target_size( out_sizes[1] = in.size(1) == 0 ? 0 : groups * weight.size(1); } - int64_t kernel_size[2]; - size_t kernel_ndim = 2; - if (weight.dim() == 3) { - kernel_size[0] = weight.size(2); - kernel_ndim = 1; - } else { - kernel_size[0] = weight.size(2); - kernel_size[1] = weight.size(3); + int64_t kernel_size[kTensorDimensionLimit]; + const size_t kernel_ndim = weight.dim() - 2; + for (const auto i : c10::irange(kernel_ndim)) { + kernel_size[i] = weight.size(i + 2); } calculate_kernel_output_sizes( in, diff --git a/kernels/test/op_convolution_test.cpp b/kernels/test/op_convolution_test.cpp index e5f8ff4f58e..3ef8b2fa05b 100644 --- a/kernels/test/op_convolution_test.cpp +++ b/kernels/test/op_convolution_test.cpp @@ -641,6 +641,308 @@ TEST_F(OpConvCorrectnessTest, InvalidInputShape) { executorch::aten::ArrayRef{output_padding, 1}, groups, out)); + + // A 3-D convolution takes one stride, padding and dilation for all spatial + // dims or one per dim, and no input has more than three spatial dims. + Tensor input_3d = tf.ones({1, 2, 3, 3, 3}); + Tensor weight_3d = tf.ones({4, 2, 2, 2, 2}); + Tensor out_3d = tf.zeros({1, 4, 2, 2, 2}); + auto conv_3d = [&](executorch::aten::ArrayRef stride_3d, + executorch::aten::ArrayRef padding_3d, + executorch::aten::ArrayRef dilation_3d) { + context_ = torch::executor::KernelRuntimeContext(); + return op_convolution_out( + input_3d, + weight_3d, + bias, + stride_3d, + padding_3d, + dilation_3d, + false, + {0}, + 1, + out_3d); + }; + const int64_t one[] = {1}; + const int64_t two[] = {2}; + const int64_t four_ones[] = {1, 1, 1, 1}; + const int64_t four_twos[] = {2, 2, 2, 2}; + // Single values apply to every spatial dim. With stride 2, padding 1 and + // dilation 2, each output sees one kernel tap per dim. The calls below differ + // only in the length of one list. + conv_3d(two, one, two); + EXPECT_TENSOR_EQ(out_3d, tf.full({1, 4, 2, 2, 2}, 2)); + ET_EXPECT_KERNEL_FAILURE(context_, conv_3d(four_twos, one, two)); + ET_EXPECT_KERNEL_FAILURE(context_, conv_3d(two, four_ones, two)); + ET_EXPECT_KERNEL_FAILURE(context_, conv_3d(two, one, four_twos)); + + Tensor input_6d = tf.ones({1, 2, 3, 3, 3, 3}); + Tensor weight_6d = tf.ones({4, 2, 2, 2, 2, 2}); + Tensor out_6d = tf.zeros({1, 4, 2, 2, 2, 2}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input_6d, weight_6d, bias, {1}, {0}, {1}, false, {0}, 1, out_6d)); + + // Like PyTorch, a kernel dim must not be empty. + Tensor narrow_input = tf.ones({1, 1, 3}); + Tensor empty_weight = tf.ones({1, 1, 0}); + Tensor wide_out = tf.zeros({1, 1, 4}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + narrow_input, + empty_weight, + bias, + {1}, + {0}, + {1}, + false, + {0}, + 1, + wide_out)); + + // Like PyTorch, an empty spatial dim is rejected even when padding would + // give it an output. + Tensor empty_input = tf.ones({1, 1, 0}); + Tensor point_weight = tf.ones({1, 1, 1}); + Tensor padded_out = tf.zeros({1, 1, 2}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + empty_input, + point_weight, + bias, + {1}, + {1}, + {1}, + false, + {0}, + 1, + padded_out)); +} + +TEST_F(OpConvCorrectnessTest, EmptyBatchOrChannelWithEmptySpatialDim) { + TensorFactory tf; + optional bias; + + // Like PyTorch, an empty spatial dim is fine when the batch or channel dim is + // empty too. + Tensor input = tf.ones({0, 1, 0}); + Tensor weight = tf.ones({1, 1, 1}); + Tensor out = tf.zeros({0, 1, 2}); + op_convolution_out(input, weight, bias, {1}, {1}, {1}, false, {0}, 1, out); + + Tensor channelless_input = tf.ones({1, 0, 0}); + Tensor channelless_weight = tf.ones({1, 0, 1}); + Tensor channelless_out = tf.zeros({1, 0, 2}); + op_convolution_out( + channelless_input, + channelless_weight, + bias, + {1}, + {1}, + {1}, + false, + {0}, + 1, + channelless_out); + + Tensor transposed_input = tf.ones({0, 2, 0}); + Tensor transposed_weight = tf.ones({2, 1, 3}); + Tensor transposed_out = tf.zeros({0, 1, 2}); + op_convolution_out( + transposed_input, + transposed_weight, + bias, + {1}, + {0}, + {1}, + true, + {0}, + 1, + transposed_out); +} + +TEST_F(OpConvCorrectnessTest, InvalidBiasLength) { + TensorFactory tf; + + Tensor input = tf.ones({1, 2, 3, 3}); + Tensor bias = tf.ones({1}); + + Tensor weight = tf.ones({4, 2, 2, 2}); + Tensor out = tf.zeros({1, 4, 2, 2}); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, weight, bias, {1}, {0}, {1}, false, {0}, 1, out)); + + Tensor transposed_weight = tf.ones({2, 4, 2, 2}); + Tensor transposed_out = tf.zeros({1, 4, 4, 4}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, + transposed_weight, + bias, + {1}, + {0}, + {1}, + true, + {0}, + 1, + transposed_out)); +} + +TEST_F(OpConvCorrectnessTest, InvalidBiasDtype) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen accepts a Bool bias"); + TensorFactory tf; + TensorFactory tf_bool; + + Tensor input = tf.ones({1, 2, 3, 3}); + Tensor weight = tf.ones({4, 2, 2, 2}); + Tensor bias = tf_bool.ones({4}); + Tensor out = tf.zeros({1, 4, 2, 2}); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, weight, bias, {1}, {0}, {1}, false, {0}, 1, out)); +} + +TEST_F(OpConvCorrectnessTest, InvalidGroups) { + TensorFactory tf; + + Tensor input = tf.ones({1, 2, 3, 3}); + Tensor weight = tf.ones({4, 2, 2, 2}); + optional bias; + Tensor out = tf.zeros({1, 4, 2, 2}); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, weight, bias, {1}, {0}, {1}, false, {0}, 0, out)); + + // A transposed bias has groups * weight.size(1) elements, so groups must be + // in range before that product is formed. Only an UndefinedBehaviorSanitizer + // build that halts on errors fails here when the order is wrong. + Tensor transposed_weight = tf.ones({2, 4, 2, 2}); + Tensor transposed_bias = tf.ones({4}); + Tensor transposed_out = tf.zeros({1, 4, 4, 4}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, + transposed_weight, + transposed_bias, + {1}, + {0}, + {1}, + true, + {0}, + std::numeric_limits::max(), + transposed_out)); +} + +TEST_F(OpConvCorrectnessTest, KernelLargerThanPaddedInput) { + TensorFactory tf; + + // Depth dilation 2 stretches the kernel depth of 2 to 3. + Tensor input = tf.ones({1, 1, 2, 2, 2}); + Tensor weight = tf.ones({1, 1, 2, 2, 2}); + optional bias; + Tensor out = tf.zeros({1, 1, 1, 1, 1}); + const int64_t stride[] = {2, 1, 1}; + const int64_t dilation[] = {2, 1, 1}; + + // Padding makes the input depth 4, so the kernel fits. + const int64_t padding[] = {1, 0, 0}; + op_convolution_out( + input, weight, bias, stride, padding, dilation, false, {0}, 1, out); + EXPECT_TENSOR_EQ(out, tf.full({1, 1, 1, 1, 1}, 4)); + + // Without padding the input depth of 2 is too small. Stride 2 must not round + // the output depth up to 1. + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, weight, bias, stride, {0}, dilation, false, {0}, 1, out)); + + // The same holds in 1-D. + Tensor input_1d = tf.ones({1, 1, 2}); + Tensor weight_1d = tf.ones({1, 1, 3}); + Tensor out_1d = tf.zeros({1, 1, 1}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input_1d, weight_1d, bias, {2}, {0}, {1}, false, {0}, 1, out_1d)); +} + +TEST_F(OpConvCorrectnessTest, HugePaddingAndDilation) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen does not reject all of these values"); + TensorFactory tf; + + Tensor input = tf.ones({1, 1, 3}); + Tensor weight = tf.ones({1, 1, 5}); + Tensor point_weight = tf.ones({1, 1, 1}); + optional bias; + const int64_t int32_max[] = {std::numeric_limits::max()}; + const int64_t huge[] = {std::numeric_limits::max() / 2 + 1}; + + // The dilated kernel size would wrap around to 1, which fits the input. + Tensor out = tf.zeros({1, 1, 3}); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, weight, bias, {1}, {0}, huge, false, {0}, 1, out)); + + // These output sizes, 2^32 + 1 and 5 - 2^32, would turn into 1 and 5 in + // 32-bit sizes. + Tensor out_1 = tf.zeros({1, 1, 1}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, + point_weight, + bias, + {1}, + int32_max, + {1}, + false, + {0}, + 1, + out_1)); + Tensor out_5 = tf.zeros({1, 1, 5}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, point_weight, bias, {1}, int32_max, {1}, true, {0}, 1, out_5)); + + // A huge stride or padding would wrap these transposed output sizes around to + // 1 and 5 in 64 bits. + Tensor input_5 = tf.ones({1, 1, 5}); + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input_5, point_weight, bias, huge, {0}, {1}, true, {0}, 1, out_1)); + const int64_t int64_max[] = {std::numeric_limits::max()}; + context_ = torch::executor::KernelRuntimeContext(); + ET_EXPECT_KERNEL_FAILURE( + context_, + op_convolution_out( + input, point_weight, bias, {1}, int64_max, {1}, true, {0}, 1, out_5)); } TEST_F(OpConvCorrectnessTest, TransposedDefaultParams) { @@ -1085,3 +1387,117 @@ TEST_F(OpConvCorrectnessTest, TransposedOutputTiles) { EXPECT_TENSOR_EQ(out, exp); } } + +namespace { +// Small integers that keep every sum exact in float. The tests pick periods +// that leave no repeats along any dim, so reading the wrong element changes the +// result. +std::vector cycling_values(size_t count, int period) { + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + values[i] = static_cast(static_cast(i) % period - period / 2); + } + return values; +} +} // namespace + +TEST_F(OpConvCorrectnessTest, Grouped3DStridePaddingDilation) { + TensorFactory tf; + // Each spatial dim has its own size, kernel size, stride, padding and + // dilation, none of which is 1 (or 0 for padding), so mixing up dims or + // dropping a factor changes the result. + const int64_t stride[] = {2, 3, 4}; + const int64_t padding[] = {1, 2, 3}; + const int64_t dilation[] = {3, 4, 2}; + auto input = tf.make({1, 4, 10, 8, 5}, cycling_values(1600, 23)); + auto weight = tf.make({2, 2, 4, 2, 3}, cycling_values(96, 11)); + auto bias = tf.make({2}, {1, -1}); + // PyTorch reference. + auto expected = + tf.make({1, 2, 2, 3, 2}, {-52, -34, -31, 64, 86, 62, -38, -87, + 81, 0, 19, 62, 56, -21, -19, -60, + 49, 26, -63, -44, 40, -173, 24, 74}); + for (const bool channels_last : {false, true}) { + SCOPED_TRACE(channels_last); + auto in = channels_last ? tf.channels_last_like(input) : input; + auto w = channels_last ? tf.channels_last_like(weight) : weight; + auto exp = channels_last ? tf.channels_last_like(expected) : expected; + auto out = tf.zeros_like(exp); + if (channels_last) { + out = tf.channels_last_like(out); + } + op_convolution_out( + in, w, bias, stride, padding, dilation, false, {0}, 2, out); + EXPECT_TENSOR_EQ(out, exp); + } + // Export gives a channels-last 5-D input a contiguous output, and the + // weight may keep its contiguous layout. + for (const Tensor& w : {weight, tf.channels_last_like(weight)}) { + auto out = tf.zeros_like(expected); + op_convolution_out( + tf.channels_last_like(input), + w, + bias, + stride, + padding, + dilation, + false, + {0}, + 2, + out); + EXPECT_TENSOR_EQ(out, expected); + } +} + +TEST_F(OpConvCorrectnessTest, GroupedTransposed3DStridePaddingDilation) { + TensorFactory tf; + // Depth has its own stride, padding, dilation and output padding, and height + // and width have their own padding and dilation. Two output channels per + // group make every weight stride matter in the channels-last layout. + const int64_t stride[] = {2, 1, 1}; + const int64_t padding[] = {3, 1, 2}; + const int64_t dilation[] = {4, 2, 3}; + const int64_t output_padding[] = {1, 0, 0}; + auto input = tf.make({1, 4, 4, 2, 3}, cycling_values(96, 13)); + auto weight = tf.make({4, 2, 2, 2, 2}, cycling_values(64, 11)); + auto bias = tf.make({4}, {1, -1, 2, -2}); + // PyTorch reference. Even output depths get only the bias. + auto expected = tf.make( + {1, 4, 6, 2, 2}, + {1, 1, 1, 1, 0, -14, 54, -31, 1, 1, 1, 1, -12, -2, -2, 2, + 1, 1, 1, 1, -13, -29, -34, -3, -1, -1, -1, -1, 6, -6, -42, -5, + -1, -1, -1, -1, -12, 0, 1, -4, -1, -1, -1, -1, 3, -1, 27, 9, + 2, 2, 2, 2, -5, -13, 3, -6, 2, 2, 2, 2, -1, 1, -11, 55, + 2, 2, 2, 2, 29, 9, 12, 4, -2, -2, -2, -2, -9, -16, 1, 5, + -2, -2, -2, -2, 1, 5, -6, -43, -2, -2, -2, -2, 28, -19, -10, -6}); + for (const bool channels_last : {false, true}) { + SCOPED_TRACE(channels_last); + auto in = channels_last ? tf.channels_last_like(input) : input; + auto w = channels_last ? tf.channels_last_like(weight) : weight; + auto exp = channels_last ? tf.channels_last_like(expected) : expected; + auto out = tf.zeros_like(exp); + if (channels_last) { + out = tf.channels_last_like(out); + } + op_convolution_out( + in, w, bias, stride, padding, dilation, true, output_padding, 2, out); + EXPECT_TENSOR_EQ(out, exp); + } + // Export gives a channels-last 5-D input a contiguous output, and the + // weight may keep its contiguous layout. + for (const Tensor& w : {weight, tf.channels_last_like(weight)}) { + auto out = tf.zeros_like(expected); + op_convolution_out( + tf.channels_last_like(input), + w, + bias, + stride, + padding, + dilation, + true, + output_padding, + 2, + out); + EXPECT_TENSOR_EQ(out, expected); + } +} diff --git a/runtime/core/exec_aten/testing_util/tensor_factory.h b/runtime/core/exec_aten/testing_util/tensor_factory.h index 562db9759a2..79cfd590f03 100644 --- a/runtime/core/exec_aten/testing_util/tensor_factory.h +++ b/runtime/core/exec_aten/testing_util/tensor_factory.h @@ -367,7 +367,8 @@ class TensorFactory { const at::Tensor& input, TensorShapeDynamism dynamism = TensorShapeDynamism::STATIC) { ET_CHECK_MSG( - input.sizes().size() == 4, "Only 4D tensors can be channels last"); + input.sizes().size() == 4 || input.sizes().size() == 5, + "Only 4D and 5D tensors can be channels last"); const std::vector sizes( input.sizes().begin(), input.sizes().end()); @@ -387,8 +388,10 @@ class TensorFactory { int32_t N = sizes[0]; int32_t C = sizes[1]; - int32_t H = sizes[2]; - int32_t W = sizes[3]; + // Depth and height are adjacent in both layouts, so a 5-D tensor is + // permuted like a 4-D one whose height covers both. + int32_t H = sizes.size() == 5 ? sizes[2] * sizes[3] : sizes[2]; + int32_t W = sizes.back(); std::vector contiguous_data( input.data_ptr(), input.data_ptr() + input.numel()); @@ -891,14 +894,18 @@ class TensorFactory { const std::vector sizes( input.sizes().begin(), input.sizes().end()); - ET_CHECK_MSG(sizes.size() == 4, "Only 4D tensors can be channels last"); + ET_CHECK_MSG( + sizes.size() == 4 || sizes.size() == 5, + "Only 4D and 5D tensors can be channels last"); ET_CHECK_MSG( is_contiguous_dim_order(input.dim_order().data(), input.dim()) == true, "Input tensor is not contiguous"); int32_t N = sizes[0]; int32_t C = sizes[1]; - int32_t H = sizes[2]; - int32_t W = sizes[3]; + // Depth and height are adjacent in both layouts, so a 5-D tensor is + // permuted like a 4-D one whose height covers both. + int32_t H = sizes.size() == 5 ? sizes[2] * sizes[3] : sizes[2]; + int32_t W = sizes.back(); std::vector contiguous_data( input.data_ptr(), input.data_ptr() + input.numel());