# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. # pyre-unsafe import ctypes import unittest from typing import Tuple import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema import torch from executorch.backends.transforms.convert_dtype_pass import I64toI32 from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend from executorch.exir import EdgeCompileConfig from torch.export import Dim, export, ExportedProgram ctypes.CDLL("libvulkan.so.1") from executorch.exir import to_edge_transform_and_lower from executorch.extension.pybindings.portable_lib import ( # @manual _load_for_executorch_from_buffer, ) from executorch.extension.pytree import tree_flatten class TestBackends(unittest.TestCase): _edge_compile_config: EdgeCompileConfig = EdgeCompileConfig( _skip_dim_order=True, # TODO(T182928844): Delegate dim order op to backend. ) def assert_outputs_equal( self, model_output, ref_output, atol=1e-03, rtol=1e-03, first_output_only=False, equal_nan=True, ): """ Helper testing function that asserts that the model output and the reference output are equal with some tolerance. Due to numerical differences between eager mode and the Vulkan's backend, we relax the detal such that default absolute tolerance is 1e-3. and default relative tolerance is 1e-3. """ # Compare the result from executor and eager mode direclty if isinstance(ref_output, tuple) or isinstance(ref_output, list): # Multiple outputs executor always returns tuple, even if there is one output self.assertTrue(len(ref_output) == len(model_output)) if first_output_only: self.assertTrue( torch.allclose( model_output[0], ref_output[0], atol=atol, rtol=rtol, equal_nan=equal_nan, ) ) else: for i in range(len(ref_output)): self.assertTrue( torch.allclose( model_output[i], ref_output[i], atol=atol, rtol=rtol, equal_nan=equal_nan, ) ) else: # If one output, eager returns tensor while executor tuple of size 1 self.assertTrue( torch.allclose( model_output[0], ref_output, atol=atol, rtol=rtol, equal_nan=equal_nan, ) ) def lower_module_and_test_output( self, model: torch.nn.Module, sample_inputs: Tuple[torch.Tensor], atol=1e-03, rtol=1e-01, dynamic_shapes=None, test_inputs=None, memory_layouts=None, first_output_only=False, ): """ Helper testing function that takes a torch.nn.Module and lowers it to Vulkan with the given sample inputs. It then runs the lowered module and compares its outputs with the outputs of the eager module. """ def run_test(memory_layout): compile_options = { "memory_layout_override": memory_layout, } # At least model should run in eager mode. model.eval() model(*sample_inputs) program: ExportedProgram = export( model, sample_inputs, dynamic_shapes=dynamic_shapes ) edge_program = to_edge_transform_and_lower( program, transform_passes=[ I64toI32(self._edge_compile_config._skip_dim_order), ], partitioner=[VulkanPartitioner(compile_options)], ) executorch_program = edge_program.to_executorch() self.assertEqual( executorch_program.executorch_program.execution_plan[0].delegates[0].id, VulkanBackend.__name__, ) executorch_module = _load_for_executorch_from_buffer( executorch_program.buffer ) inputs_flattened, _ = tree_flatten(sample_inputs) model_output = executorch_module.run_method( "forward", tuple(inputs_flattened) ) ref_output = model(*sample_inputs) self.assert_outputs_equal( model_output, ref_output, atol=atol, rtol=rtol, first_output_only=first_output_only, ) if test_inputs is not None: for test_input in test_inputs: test_inputs_flattened, _ = tree_flatten(test_input) model_output = executorch_module.run_method( "forward", tuple(test_inputs_flattened) ) ref_output = model(*test_input) self.assert_outputs_equal( model_output, ref_output, atol=atol, rtol=rtol, first_output_only=first_output_only, ) memory_layouts_to_test = [ vk_graph_schema.VkMemoryLayout.TENSOR_WIDTH_PACKED, vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED, ] if memory_layouts is not None: memory_layouts_to_test = memory_layouts for memory_layout in memory_layouts_to_test: run_test(memory_layout) def test_vulkan_backend_add(self): # This test is the simplest test by manually lowering some submodules, we can use paritioner # for auto detecting lowerable parts. class AddModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y, w): z = x + y z = z + x z = z + x z = z + w z = w + z z = z + 3 # test scalar broadcasting return z add_module = AddModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 1), dtype=torch.float32), # test broadcasting ) self.lower_module_and_test_output(add_module, sample_inputs) sample_inputs = ( torch.rand(size=(4, 5, 2, 3), dtype=torch.float32), torch.rand(size=(4, 5, 2, 3), dtype=torch.float32), torch.rand( size=(2, 3), dtype=torch.float32 ), # test broadcasting on packed dim ) self.lower_module_and_test_output(add_module, sample_inputs) def test_vulkan_backend_add_int(self): class AddIntModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): z = x + y return z add_int_module = AddIntModule() sample_inputs = ( torch.randint(low=-100, high=100, size=(2, 3), dtype=torch.int32), torch.randint(low=-100, high=100, size=(2, 3), dtype=torch.int32), ) self.lower_module_and_test_output(add_int_module, sample_inputs) def test_vulkan_backend_zero_dim_tensor(self): class ZeroDimModule(torch.nn.Module): def __init__(self): super().__init__() self.zero = torch.full([], 1.3, dtype=torch.float32) def forward(self, x): return x + self.zero internal_data_module = ZeroDimModule() sample_inputs = (torch.rand(size=(2, 3), dtype=torch.float32),) self.lower_module_and_test_output(internal_data_module, sample_inputs) def test_vulkan_backend_internal_data(self): class InternalDataModule(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.rand(size=(2, 3), dtype=torch.float32) def forward(self, x, y): inter1 = torch.add(x, y, alpha=2) inter2 = torch.add(x, y, alpha=3.14) inter3 = inter1 * self.weight inter4 = inter2 * self.weight return inter4 - inter3 internal_data_module = InternalDataModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), ) self.lower_module_and_test_output(internal_data_module, sample_inputs) def test_vulkan_backend_sub(self): class SubModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): z = torch.sub(x, y, alpha=2) z = torch.sub(z, x, alpha=3.14) z = z - x return z sub_module = SubModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), ) self.lower_module_and_test_output(sub_module, sample_inputs) def test_vulkan_backend_mul(self): class MulModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): z = x * y z = z * x z = z * x return z mul_module = MulModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), ) self.lower_module_and_test_output(mul_module, sample_inputs) def test_vulkan_backend_div(self): class DivModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): z = x / y z = z / x z = z / x return z div_module = DivModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), ) self.lower_module_and_test_output(div_module, sample_inputs) def test_vulkan_backend_arithmetic(self): class ArithmeticModule(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.rand(size=(2, 3), dtype=torch.float32) def forward(self, x, y): z = x + y z = z - x z = z / x z = z * self.weight return z arithmetic_module = ArithmeticModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), ) self.lower_module_and_test_output(arithmetic_module, sample_inputs) def test_vulkan_backend_floor_div(self): class FloorDivModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): z = x // y return z floor_div_module = FloorDivModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32) * 10.0, torch.rand(size=(2, 3), dtype=torch.float32) + 1.0, ) # absolute tolerance is 1 because of flooring self.lower_module_and_test_output( floor_div_module, sample_inputs, atol=1.0 + 1e-03 ) def test_vulkan_backend_pow(self): class PowModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): z = torch.pow(x, y) return z pow_module = PowModule() sample_inputs = ( torch.rand(size=(2, 3), dtype=torch.float32), torch.rand(size=(2, 3), dtype=torch.float32), ) self.lower_module_and_test_output(pow_module, sample_inputs) def lower_unary_module_and_test_output(self, module): batch = Dim("batch", max=8) sample_inputs = (torch.randn(8, 16, 96, 92),) dynamic_shapes = {"x": {0: batch}} test_inputs = [ (torch.randn(3, 14, 15, 92),), (torch.randn(6, 5, 35, 89),), (torch.randn(7, 9, 32, 38),), ] self.lower_module_and_test_output( module, sample_inputs, dynamic_shapes=dynamic_shapes, test_inputs=test_inputs, ) def test_vulkan_backend_clamp(self): class ClampModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.clamp(x, min=-3.14) self.lower_unary_module_and_test_output(ClampModule()) def test_vulkan_backend_clamp_int(self): class ClampModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.clamp(x, min=-3) sample_inputs = ( torch.randint(low=-100, high=100, size=(5, 5), dtype=torch.int32), ) self.lower_module_and_test_output(ClampModule(), sample_inputs) def test_vulkan_backend_clamp_int64(self): class ClampModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.clamp(x, min=-3) sample_inputs = ( torch.randint(low=-100, high=100, size=(5, 5), dtype=torch.int64), ) self.lower_module_and_test_output(ClampModule(), sample_inputs) def test_vulkan_backend_cos(self): class CosModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.cos(x) self.lower_unary_module_and_test_output(CosModule()) def test_vulkan_backend_hardtanh(self): class HardTanHModule(torch.nn.Module): def __init__(self): super().__init__() self.tanh = torch.nn.Hardtanh(min_val=-3.14, max_val=6.28) def forward(self, x): return self.tanh(x) self.lower_unary_module_and_test_output(HardTanHModule()) def test_vulkan_backend_exp(self): class ExpModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.exp(x) self.lower_unary_module_and_test_output(ExpModule()) def test_vulkan_backend_neg(self): class NegModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.neg(x) self.lower_unary_module_and_test_output(NegModule()) def test_vulkan_backend_sin(self): class SinModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.sin(x) self.lower_unary_module_and_test_output(SinModule()) def test_vulkan_backend_relu(self): class ReLUModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.relu(x) self.lower_unary_module_and_test_output(ReLUModule()) def test_vulkan_backend_sqrt(self): class SqrtModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.sqrt(x) self.lower_unary_module_and_test_output(SqrtModule()) def test_vulkan_backend_hardshrink(self): class HardshrinkModule(torch.nn.Module): def __init__(self): super().__init__() self.hardshrink = torch.nn.Hardshrink(lambd=0.3) def forward(self, x): return self.hardshrink(x) self.lower_unary_module_and_test_output(HardshrinkModule()) def test_vulkan_backend_max_pool2d(self): class MaxPool2dModule(torch.nn.Module): def __init__(self): super().__init__() self.max_pool = torch.nn.MaxPool2d( kernel_size=(2, 3), stride=(1, 1), padding=0, dilation=1, ceil_mode=False, return_indices=True, ) def forward(self, x): return self.max_pool(x) max_pool2d_module = MaxPool2dModule() sample_inputs = (torch.randn(5, 13, 55, 68),) batch = Dim("batch", max=8) dynamic_shapes = {"x": {0: batch}} test_inputs = [ (torch.randn(3, 14, 15, 9),), (torch.randn(1, 1, 4, 6),), (torch.randn(5, 10, 50, 40),), ] self.lower_module_and_test_output( max_pool2d_module, sample_inputs, dynamic_shapes=dynamic_shapes, test_inputs=test_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], first_output_only=True, ) def test_vulkan_backend_avg_pool2d(self): class AvgPool2dModule(torch.nn.Module): def __init__(self): super().__init__() self.avg_pool = torch.nn.AvgPool2d( kernel_size=(4, 4), stride=(4, 4), padding=(0, 0), ceil_mode=True, count_include_pad=True, divisor_override=None, ) def forward(self, x): return self.avg_pool(x) avg_pool2d_module = AvgPool2dModule() sample_inputs = (torch.randn(5, 13, 55, 68),) batch = Dim("batch", max=8) dynamic_shapes = {"x": {0: batch}} test_inputs = [ (torch.randn(3, 14, 15, 9),), (torch.randn(1, 1, 4, 6),), (torch.randn(5, 10, 50, 40),), ] self.lower_module_and_test_output( avg_pool2d_module, sample_inputs, dynamic_shapes=dynamic_shapes, test_inputs=test_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_abs(self): class AbsModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.abs(x) self.lower_unary_module_and_test_output(AbsModule()) def test_vulkan_backend_sigmoid(self): class SigmoidModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.sigmoid(x) self.lower_unary_module_and_test_output(SigmoidModule()) def test_vulkan_backend_tanh(self): class TanhModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.tanh(x) self.lower_unary_module_and_test_output(TanhModule()) def test_vulkan_backend_linear(self): class LinearModule(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(128, 64, bias=False) def forward(self, x): return self.linear(x) module = LinearModule() sample_inputs = (torch.rand(size=(32, 128), dtype=torch.float32),) batch = Dim("batch", max=32) dynamic_shapes = {"x": {0: batch}} test_inputs = [ (torch.rand(15, 128),), (torch.rand(6, 128),), (torch.rand(30, 128),), (torch.rand(20, 128),), (torch.rand(19, 128),), ] self.lower_module_and_test_output( module, sample_inputs, dynamic_shapes=dynamic_shapes, test_inputs=test_inputs, ) def test_vulkan_backend_partial(self): class SimpleModel(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(10, 10) self.offset_1 = torch.rand(size=(2, 10), dtype=torch.float32) self.offset_2 = torch.rand(size=(2, 10), dtype=torch.float32) def forward(self, x): return self.linear(x + self.offset_1) - self.offset_2 model = SimpleModel() sample_inputs = (torch.rand(size=(2, 10), dtype=torch.float32),) self.lower_module_and_test_output(model, sample_inputs) def test_vulkan_backend_partial_dynamic_shapes(self): class SimpleModel(torch.nn.Module): def __init__(self): super().__init__() self.branch1 = torch.nn.Sequential( torch.nn.Linear(64, 64), torch.nn.ReLU() ) self.branch2 = torch.nn.Sequential( torch.nn.Linear(128, 64), torch.nn.ReLU() ) self.buffer_1 = torch.ones((1, 64)) * 0.5 self.buffer_2 = torch.ones((1, 64)) * 1.4 def forward(self, x1, x2): out1 = self.branch1(x1) out2 = self.branch2(x2) return (out1 + self.buffer_1 + out2) * self.buffer_2 model = SimpleModel() sample_inputs = (torch.randn(32, 64), torch.randn(32, 128)) batch = Dim("batch", max=32) dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}} test_inputs = [ (torch.randn(15, 64), torch.randn(15, 128)), (torch.randn(6, 64), torch.randn(6, 128)), (torch.randn(30, 64), torch.randn(30, 128)), (torch.randn(20, 64), torch.randn(20, 128)), (torch.randn(19, 64), torch.randn(19, 128)), ] self.lower_module_and_test_output( model, sample_inputs, dynamic_shapes=dynamic_shapes, test_inputs=test_inputs ) def test_vulkan_backend_matmul(self): class MatMulModule(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.ones(size=(63, 22), dtype=torch.float32) def forward(self, x): return torch.matmul(x, self.weight) module = MatMulModule() sample_inputs = (torch.ones(size=(31, 63), dtype=torch.float32),) self.lower_module_and_test_output(module, sample_inputs) def test_vulkan_backend_bmm(self): class BMMModule(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.randn(size=(4, 4, 5), dtype=torch.float32) def forward(self, x): return torch.bmm(x, self.weight) module = BMMModule() sample_inputs = (torch.randn(size=(4, 3, 4), dtype=torch.float32),) self.lower_module_and_test_output(module, sample_inputs) @unittest.skip( "Reduce shader does not support multiple reduction axes at the moment" ) def test_vulkan_backend_sum_dim_list(self): class SumModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): x = torch.sum(x, (0, -1), keepdim=True) x = torch.sum(x, 2, keepdim=False) return x module = SumModule() sample_inputs = (torch.ones(size=(3, 2, 7, 5), dtype=torch.float32),) self.lower_module_and_test_output( module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) @unittest.skip( "Reduce shader does not support multiple reduction axes at the moment" ) def test_vulkan_backend_sum(self): class SumModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): x = torch.sum(x, (), keepdim=True) x = torch.sum(x) return x module = SumModule() sample_inputs = (torch.rand(size=(3, 2, 7, 5), dtype=torch.float32),) self.lower_module_and_test_output( module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv2d(self): class Conv2dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d( in_channels=6, out_channels=8, kernel_size=(3, 3), padding=(2, 3), stride=(1, 2), dilation=1, groups=1, bias=True, ) def forward(self, x): return self.conv(x) conv2d_module = Conv2dModule() sample_inputs = (torch.randn(size=(1, 6, 40, 50), dtype=torch.float32),) self.lower_module_and_test_output( conv2d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv_transpose2d(self): class ConvTranspose2dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.ConvTranspose2d( in_channels=6, out_channels=8, kernel_size=(3, 3), padding=(2, 3), stride=(1, 2), output_padding=(0, 1), dilation=1, groups=1, bias=True, ) def forward(self, x): return self.conv(x) conv_transpose2d_module = ConvTranspose2dModule() sample_inputs = (torch.randn(size=(1, 6, 40, 50), dtype=torch.float32),) self.lower_module_and_test_output( conv_transpose2d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv2d_dw(self): class Conv2dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d( in_channels=8, out_channels=8, kernel_size=3, padding=1, groups=8, bias=True, ) def forward(self, x): return self.conv(x) conv2d_module = Conv2dModule() sample_inputs = (torch.randn(size=(1, 8, 72, 96), dtype=torch.float32),) self.lower_module_and_test_output( conv2d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv2d_pw(self): class Conv2dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d( in_channels=8, out_channels=8, kernel_size=1, padding=1, groups=1, bias=True, ) def forward(self, x): return self.conv(x) conv2d_module = Conv2dModule() sample_inputs = (torch.randn(size=(1, 8, 72, 96), dtype=torch.float32),) self.lower_module_and_test_output( conv2d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv2d_bias_false(self): class Conv2dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d( in_channels=6, out_channels=8, kernel_size=(3, 3), padding=(2, 3), stride=(1, 2), dilation=1, groups=1, bias=False, ) def forward(self, x): return self.conv(x) conv2d_module = Conv2dModule() sample_inputs = (torch.randn(size=(1, 6, 40, 50), dtype=torch.float32),) self.lower_module_and_test_output( conv2d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv1d(self): class Conv1dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv1d( in_channels=20, out_channels=10, kernel_size=6, stride=5, padding=5, dilation=3, groups=5, bias=True, ) def forward(self, x): return self.conv(x) conv1d_module = Conv1dModule() sample_inputs = (torch.randn(size=(3, 20, 30), dtype=torch.float32),) self.lower_module_and_test_output( conv1d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv1d_bias_false(self): class Conv1dModule(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv1d( in_channels=6, out_channels=6, kernel_size=3, groups=6, bias=False, ) def forward(self, x): return self.conv(x) conv1d_module = Conv1dModule() sample_inputs = (torch.randn(size=(1, 6, 7), dtype=torch.float32),) self.lower_module_and_test_output( conv1d_module, sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_native_layer_norm(self): class NativeLayerNormModule(torch.nn.Module): def __init__(self): super().__init__() self.layer_norm = torch.nn.LayerNorm(5) def forward(self, x): return self.layer_norm(x) sample_inputs = (torch.randn(size=(3, 4, 5), dtype=torch.float32),) self.lower_module_and_test_output( NativeLayerNormModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_batch_norm(self): class BatchNormModule(torch.nn.Module): def __init__(self): super().__init__() self.bn = torch.nn.BatchNorm2d(num_features=3) def forward(self, x): return self.bn(x) sample_inputs = (torch.randn(size=(4, 3, 2, 5), dtype=torch.float32),) self.lower_module_and_test_output( BatchNormModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_full(self): class FullModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.full(x.shape, 42.0) class ZerosModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.zeros(x.shape) class OnesModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.ones(x.shape) sample_inputs = (torch.randn(size=(2, 3, 4, 5), dtype=torch.float32),) self.lower_module_and_test_output( FullModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( ZerosModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( OnesModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_full_like(self): class FullLikeModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.full_like(x, 42.0) class ZerosLikeModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.zeros_like(x) class OnesLikeModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.ones_like(x) sample_inputs = (torch.randn(size=(2, 3, 4, 5), dtype=torch.float32),) self.lower_module_and_test_output( FullLikeModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( ZerosLikeModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( OnesLikeModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_upsample_nearest2d(self): class UpsampleNearest2d(torch.nn.Module): def __init__(self): super().__init__() self.upsample = torch.nn.Upsample(scale_factor=2, mode="nearest") def forward(self, x): return self.upsample(x) sample_inputs = (torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2),) self.lower_module_and_test_output( UpsampleNearest2d(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_minimum(self): class MinimumModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): return torch.minimum(x, y) sample_inputs = ( torch.rand(size=(3, 5, 6, 4), dtype=torch.float32), torch.rand(size=(6, 4), dtype=torch.float32), ) self.lower_module_and_test_output( MinimumModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_reshape(self): class ReshapeModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.reshape(x, [-1, x.size(-1)]) sample_inputs = (torch.randn(size=(5, 3, 4), dtype=torch.float32),) self.lower_module_and_test_output( ReshapeModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_view(self): class ViewModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return x.view([-1, x.size(-1)]) sample_inputs = (torch.randn(size=(3, 2, 3, 4), dtype=torch.float32),) self.lower_module_and_test_output( ViewModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_view_int(self): class ViewModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return x.view([-1, x.size(-1)]) sample_inputs = (torch.randint(size=(3, 6, 2, 7), high=100, dtype=torch.int32),) self.lower_module_and_test_output( ViewModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_unsqueeze(self): class UnsqueezeModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): x = torch.unsqueeze(x, 1) x = torch.unsqueeze(x, 0) return x sample_inputs = (torch.randn(size=(3,), dtype=torch.float32),) self.lower_module_and_test_output( UnsqueezeModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_squeeze(self): class SqueezeModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.squeeze(x, 0) sample_inputs = (torch.randn(size=(1, 2, 2, 1), dtype=torch.float32),) self.lower_module_and_test_output( SqueezeModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_select(self): class SelectModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return x[0][3] sample_inputs = (torch.randn(size=(3, 6, 2, 7), dtype=torch.float32),) self.lower_module_and_test_output( SelectModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_permute_copy(self): class PermuteModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.permute(x, [3, 0, 2, 1]) sample_inputs = (torch.randn(size=(3, 6, 2, 7), dtype=torch.float32),) self.lower_module_and_test_output( PermuteModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_permute_copy_int(self): class PermuteModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.permute(x, [3, 0, 2, 1]) sample_inputs = (torch.randint(size=(3, 6, 2, 7), high=100, dtype=torch.int32),) self.lower_module_and_test_output( PermuteModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_cat(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y, z, w): return torch.cat([x, y, z, w], dim=1) sample_inputs = ( torch.randn(size=(3, 6, 2, 7), dtype=torch.float32), torch.randn(size=(3, 1, 2, 7), dtype=torch.float32), torch.randn(size=(3, 9, 2, 7), dtype=torch.float32), torch.randn(size=(3, 3, 2, 7), dtype=torch.float32), ) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_cat_with_zero_size(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y, z, w): return torch.cat([x, y, z, w], dim=1) sample_inputs = ( torch.randn(size=(3, 6, 2, 7), dtype=torch.float32), torch.randn(size=(3, 0, 2, 7), dtype=torch.float32), torch.randn(size=(3, 0, 2, 7), dtype=torch.float32), torch.randn(size=(3, 3, 2, 7), dtype=torch.float32), ) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_slice(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return x[:, 2:9:2, :] sample_inputs = (torch.randn(size=(3, 13, 7, 3), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_split_with_sizes(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.split(x, (3, 6, 1, 3), dim=1) sample_inputs = (torch.randn(size=(3, 13, 7, 3), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_split_tensor(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.tensor_split(x, 2, dim=1) sample_inputs = (torch.randn(size=(3, 14, 7, 3), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_clone(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.clone(x) sample_inputs = (torch.randn(size=(3, 14, 7, 3), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_constant_pad_nd(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.nn.functional.pad(x, (1, 2, 3, 4, 5, 6), "constant", 24.2) sample_inputs = (torch.randn(size=(3, 7, 5, 11), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_repeat(self): class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return x.repeat([2, 3, 1, 2]) sample_inputs = (torch.randn(size=(3, 7, 5, 9), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_t_default(self): # aten.permute_copy.default is not enabled yet in partitioner class TestModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): # torch.t is actually exported as aten::permute. return torch.t(x) sample_inputs = (torch.randn(size=(3, 14), dtype=torch.float32),) self.lower_module_and_test_output( TestModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) @unittest.skip( "Softmax shader with shared memory does not work with swiftshader due to potential swiftshader bug" ) def test_vulkan_backend_softmax(self): class SoftmaxModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): x = x.softmax(dim=0) x = x.softmax(dim=1) x = x.softmax(dim=2) return x sample_inputs = (torch.randn(size=(3, 2, 7), dtype=torch.float32),) self.lower_module_and_test_output( SoftmaxModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) @unittest.skip( "Softmax shader with shared memory does not work with swiftshader due to potential swiftshader bug" ) def test_vulkan_backend_logsoftmax(self): class LogSoftmaxModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): x = x.log_softmax(dim=0) x = x.log_softmax(dim=1) x = x.log_softmax(dim=2) return x sample_inputs = (torch.randn(size=(3, 2, 7), dtype=torch.float32),) self.lower_module_and_test_output( LogSoftmaxModule(), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_gelu(self): class GeluModule(torch.nn.Module): def __init__(self): super().__init__() self.gelu = torch.nn.GELU(approximate="tanh") def forward(self, x): return self.gelu(x) self.lower_unary_module_and_test_output(GeluModule()) @unittest.skip( "Reduce shader does not support multiple reduction axes at the moment" ) def test_vulkan_backend_mean(self): class MeanModule(torch.nn.Module): def __init__(self, dims, keepdim=True): super().__init__() self.dims = dims self.keepdim = keepdim def forward(self, x): return torch.mean(x, self.dims, keepdim=self.keepdim) sample_inputs = ( torch.arange(end=2 * 3 * 2 * 5, dtype=torch.float32).reshape(2, 3, 2, 5), ) self.lower_module_and_test_output( MeanModule(dims=[-1, -2]), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( MeanModule(dims=[1]), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( MeanModule(dims=[0, 1, 2, 3]), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( MeanModule(dims=[-1, -2], keepdim=False), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( MeanModule(dims=[1], keepdim=False), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_index_select_int(self): class IndexSelectModule(torch.nn.Module): def __init__(self, dim, indices): super().__init__() self.dim = dim self.index = torch.tensor(indices) def forward(self, x): return torch.index_select(x, self.dim, self.index) sample_inputs = (torch.arange(96).reshape(2, 8, 2, 3),) self.lower_module_and_test_output( IndexSelectModule(dim=1, indices=[2, 3, 5, 6, 7]), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_index_select(self): class IndexSelectModule(torch.nn.Module): def __init__(self, dim, indices): super().__init__() self.dim = dim self.index = torch.tensor(indices) def forward(self, x): return torch.index_select(x, self.dim, self.index) sample_inputs = (torch.arange(144).reshape(12, 1, 3, 4).float(),) self.lower_module_and_test_output( IndexSelectModule(dim=0, indices=[1, 3, 5, 7, 8, 9, 10, 11, 2, 3]), sample_inputs, memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_arange_int(self): class ArangeModule(torch.nn.Module): def __init__(self, input): super().__init__() self.input = input def forward(self, x): return torch.arange(*self.input, dtype=torch.int32) # `torch.arange` could take one, two or three arguments as input. # If only one argument is provided, it will be interpreted as `end`. # If two arguments are provided, the first one will be interpreted as `start` # and the second one will be interpreted as `end`. # If three arguments are provided, the first one will be interpreted as `start`, # the second one will be interpreted as `end` and the third one will be # interpreted as `step`. inputs = [ [1], [-3, 5], [1, 11, 2], [12, 1, -2], ] for i in inputs: self.lower_module_and_test_output( ArangeModule(i), (torch.randn(size=(1,), dtype=torch.float32),), # dummy input memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_arange_float(self): class ArangeModule(torch.nn.Module): def __init__(self, input): super().__init__() self.input = input def forward(self, x): return torch.arange(*self.input) inputs = [ [1.5], [-3, 5.0], [1.0, 11, 2], [12, 1, -2.0], ] for i in inputs: self.lower_module_and_test_output( ArangeModule(i), (torch.randn(size=(1,), dtype=torch.float32),), # dummy input memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_arange_int64(self): class ArangeModule(torch.nn.Module): def __init__(self, input): super().__init__() self.input = input def forward(self, x): return torch.arange(*self.input) inputs = [ [1], [-3, 5], [1, 11, 2], [12, 1, -2], [1.5], [-3, 5.0], [1.0, 11, 2], [12, 1, -2.0], ] for i in inputs: self.lower_module_and_test_output( ArangeModule(i), (torch.randn(size=(1,), dtype=torch.float32),), # dummy input memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) self.lower_module_and_test_output( ArangeModule(i), (torch.randint(low=-100, high=100, size=(5, 5)),), # dummy input memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_embedding_1d(self): class EmbeddingModule(torch.nn.Module): def __init__(self, embedding): super().__init__() self.embedding = embedding def forward(self, x): return self.embedding(x) self.lower_module_and_test_output( EmbeddingModule(torch.nn.Embedding(5, 4)), (torch.tensor([0, 1, 0, 4, 2, 0]),), memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_embedding_2d(self): class EmbeddingModule(torch.nn.Module): def __init__(self, embedding): super().__init__() self.embedding = embedding def forward(self, x): return self.embedding(x) self.lower_module_and_test_output( EmbeddingModule(torch.nn.Embedding(5, 4)), (torch.tensor([[0, 1, 0], [4, 2, 0]]),), memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_embedding_3d(self): class EmbeddingModule(torch.nn.Module): def __init__(self, embedding): super().__init__() self.embedding = embedding def forward(self, x): return self.embedding(x) self.lower_module_and_test_output( EmbeddingModule(torch.nn.Embedding(5, 4)), (torch.tensor([[[0, 1], [0, 1]], [[4, 2], [3, 3]]]),), memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_flip(self): class FlipModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.flip(x, [0, 1, 2, 3]) self.lower_module_and_test_output( FlipModule(), (torch.arange(48).reshape(2, 3, 4, 2),), memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_conv_with_clamp(self): class ConvWithClampModule(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.randn(6, 8, 3, 3) self.bias = torch.randn(8) self.stride = (1, 2) self.padding = (2, 3) self.dilation = (1, 1) self.transposed = True self.output_padding = (0, 1) self.groups = 1 self.output_min = 0 self.output_max = 10 def forward(self, x): return torch.ops.et_vk.conv_with_clamp( x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.transposed, self.output_padding, self.groups, self.output_min, self.output_max, ) self.lower_module_and_test_output( ConvWithClampModule(), (torch.randn(size=(1, 6, 40, 50), dtype=torch.float32),), memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], ) def test_vulkan_backend_grid_priors(self): class GridPriorsModule(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.ops.et_vk.grid_priors( x, stride=8, offset=0.5, ) self.lower_module_and_test_output( GridPriorsModule(), (torch.rand(size=[1, 5, 2, 3]),), memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED], )