# Owner(s): ["oncall: mobile"] import io import itertools import unittest from hypothesis import assume, given, strategies as st import torch import torch.backends.xnnpack import torch.testing._internal.hypothesis_utils as hu from torch.nn import functional as F from torch.testing import FileCheck from torch.testing._internal.common_utils import ( IS_FBCODE, run_tests, slowTest, TEST_WITH_TSAN, TestCase, ) from torch.utils.mobile_optimizer import optimize_for_mobile @unittest.skipUnless( torch.backends.xnnpack.enabled, " XNNPACK must be enabled for these tests." " Please build with USE_XNNPACK=1.", ) @unittest.skipIf( TEST_WITH_TSAN, "TSAN fails with XNNPACK. Does not seem to have a good reason for failures.", ) class TestXNNPACKOps(TestCase): @unittest.skip( "Fails on some platforms, see https://github.com/pytorch/pytorch/issues/73488" ) @given( batch_size=st.integers(0, 3), data_shape=hu.array_shapes(1, 3, 2, 64), weight_output_dim=st.integers(2, 64), use_bias=st.booleans(), ) def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias): data_shape = [batch_size] + list(data_shape) input_data = torch.rand(data_shape) weight = torch.rand((weight_output_dim, data_shape[-1])) if use_bias: bias = torch.rand(weight_output_dim) else: bias = None ref_result = F.linear(input_data, weight, bias) packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias) output_linearprepacked = torch.ops.prepacked.linear_clamp_run( input_data, packed_weight_bias ) torch.testing.assert_close( ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3 ) @given( input_size=st.integers(2, 32), weight_output_dim=st.integers(2, 64), use_bias=st.booleans(), ) def test_linear_1d_input(self, input_size, weight_output_dim, use_bias): input_data = torch.rand(input_size) weight = torch.rand((weight_output_dim, input_data.shape[-1])) if use_bias: bias = torch.rand(weight_output_dim) else: bias = None ref_result = F.linear(input_data, weight, bias) packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias) output_linearprepacked = torch.ops.prepacked.linear_clamp_run( input_data, packed_weight_bias ) torch.testing.assert_close( ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3 ) @given( batch_size=st.integers(0, 3), input_channels_per_group=st.integers(1, 32), height=st.integers(5, 64), width=st.integers(5, 64), output_channels_per_group=st.integers(1, 32), groups=st.integers(1, 16), kernel_h=st.integers(1, 7), kernel_w=st.integers(1, 7), stride_h=st.integers(1, 2), stride_w=st.integers(1, 2), pad_h=st.integers(0, 2), pad_w=st.integers(0, 2), dilation=st.integers(1, 2), use_bias=st.booleans(), format=st.sampled_from( [None, torch.preserve_format, torch.contiguous_format, torch.channels_last] ), ) def test_conv2d( self, batch_size, input_channels_per_group, height, width, output_channels_per_group, groups, kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, dilation, use_bias, format, ): input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups kernels = (kernel_h, kernel_w) strides = (stride_h, stride_w) paddings = (pad_h, pad_w) dilations = (dilation, dilation) assume(height + 2 * paddings[0] >= dilations[0] * (kernels[0] - 1) + 1) assume(width + 2 * paddings[1] >= dilations[1] * (kernels[1] - 1) + 1) input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) weight = torch.rand( (output_channels, input_channels_per_group, kernel_h, kernel_w) ) bias = None if use_bias: bias = torch.rand(output_channels) ref_result = F.conv2d( input_data, weight, bias, strides, paddings, dilations, groups ) packed_weight_bias = torch.ops.prepacked.conv2d_clamp_prepack( weight, bias, strides, paddings, dilations, groups ) xnnpack_result = torch.ops.prepacked.conv2d_clamp_run( input_data, packed_weight_bias ) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) @given( batch_size=st.integers(1, 3), input_channels_per_group=st.integers(1, 32), height=st.integers(5, 64), width=st.integers(5, 64), output_channels_per_group=st.integers(1, 32), groups=st.integers(1, 16), kernel_h=st.integers(1, 7), kernel_w=st.integers(1, 7), stride_h=st.integers(1, 2), stride_w=st.integers(1, 2), pad_h=st.integers(0, 2), pad_w=st.integers(0, 2), output_pad_h=st.integers(0, 2), output_pad_w=st.integers(0, 2), dilation=st.integers(1, 2), use_bias=st.booleans(), format=st.sampled_from( [None, torch.preserve_format, torch.contiguous_format, torch.channels_last] ), ) def test_conv2d_transpose( self, batch_size, input_channels_per_group, height, width, output_channels_per_group, groups, kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, output_pad_h, output_pad_w, dilation, use_bias, format, ): input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups kernels = (kernel_h, kernel_w) strides = (stride_h, stride_w) paddings = (pad_h, pad_w) output_paddings = (output_pad_h, output_pad_w) dilations = (dilation, dilation) assume(height + 2 * paddings[0] >= dilations[0] * (kernels[0] - 1) + 1) assume(width + 2 * paddings[1] >= dilations[1] * (kernels[1] - 1) + 1) assume((output_pad_h < stride_h) and (output_pad_h < dilation)) assume((output_pad_w < stride_w) and (output_pad_w < dilation)) input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) weight = torch.rand( (input_channels, output_channels_per_group, kernel_h, kernel_w) ) bias = None if use_bias: bias = torch.rand(output_channels) # Note that groups/dilation is in reverse order from conv2d ref_result = F.conv_transpose2d( input_data, weight, bias, strides, paddings, output_paddings, groups, dilation, ) packed_weight_bias = torch.ops.prepacked.conv2d_transpose_clamp_prepack( weight, bias, strides, paddings, output_paddings, dilations, groups ) xnnpack_result = torch.ops.prepacked.conv2d_transpose_clamp_run( input_data, packed_weight_bias ) torch.testing.assert_close( ref_result.contiguous(), xnnpack_result.contiguous(), rtol=1e-2, atol=1e-3 ) @unittest.skipUnless( torch.backends.xnnpack.enabled, " XNNPACK must be enabled for these tests." " Please build with USE_XNNPACK=1.", ) @unittest.skipIf( TEST_WITH_TSAN, "TSAN fails with XNNPACK. Does not seem to have a good reason for failures.", ) class TestXNNPACKSerDes(TestCase): @unittest.skip( "Fails on some platforms, see https://github.com/pytorch/pytorch/issues/73488" ) @given( batch_size=st.integers(0, 3), data_shape=hu.array_shapes(1, 3, 2, 64), weight_output_dim=st.integers(2, 64), use_bias=st.booleans(), ) def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias): class Linear(torch.nn.Module): def __init__(self, weight, bias=None): super().__init__() self.weight = weight self.bias = bias def forward(self, x): return F.linear(x, self.weight, self.bias) class LinearPrePacked(torch.nn.Module): def __init__(self, weight, bias=None): super().__init__() self.packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack( weight, bias ) def forward(self, x): return torch.ops.prepacked.linear_clamp_run(x, self.packed_weight_bias) data_shape = [batch_size] + list(data_shape) weight = torch.rand((weight_output_dim, data_shape[-1])) if use_bias: bias = torch.rand(weight_output_dim) else: bias = None scripted_linear = torch.jit.script(Linear(weight, bias)) scripted_linear_clamp_prepacked = torch.jit.script( LinearPrePacked(weight, bias) ) input_data = torch.rand(data_shape) ref_result = scripted_linear(input_data) output_linearprepacked = scripted_linear_clamp_prepacked(input_data) torch.testing.assert_close( ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3 ) # Serialize the modules and then deserialize input_data = torch.rand(data_shape) buffer = io.BytesIO() torch.jit.save(scripted_linear, buffer) buffer.seek(0) deserialized_linear = torch.jit.load(buffer) buffer = io.BytesIO() torch.jit.save(scripted_linear_clamp_prepacked, buffer) buffer.seek(0) deserialized_linear_clamp_prepacked = torch.jit.load(buffer) ref_result = deserialized_linear(input_data) output_linearprepacked = deserialized_linear_clamp_prepacked(input_data) torch.testing.assert_close( ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3 ) @given( batch_size=st.integers(0, 3), input_channels_per_group=st.integers(1, 32), height=st.integers(5, 64), width=st.integers(5, 64), output_channels_per_group=st.integers(1, 32), groups=st.integers(1, 16), kernel_h=st.integers(1, 7), kernel_w=st.integers(1, 7), stride_h=st.integers(1, 2), stride_w=st.integers(1, 2), pad_h=st.integers(0, 2), pad_w=st.integers(0, 2), dilation=st.integers(1, 2), use_bias=st.booleans(), format=st.sampled_from( [None, torch.preserve_format, torch.contiguous_format, torch.channels_last] ), ) def test_conv2d( self, batch_size, input_channels_per_group, height, width, output_channels_per_group, groups, kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, dilation, use_bias, format, ): class Conv2D(torch.nn.Module): def __init__(self, weight, bias, strides, paddings, dilations, groups): super().__init__() self.weight = weight self.bias = bias self.strides = strides self.paddings = paddings self.dilations = dilations self.groups = groups def forward(self, x): return F.conv2d( x, self.weight, self.bias, self.strides, self.paddings, self.dilations, self.groups, ) class Conv2DPrePacked(torch.nn.Module): def __init__(self, weight, bias, strides, paddings, dilations, groups): super().__init__() self.packed_weight_bias = torch.ops.prepacked.conv2d_clamp_prepack( weight, bias, strides, paddings, dilations, groups ) def forward(self, x): return torch.ops.prepacked.conv2d_clamp_run(x, self.packed_weight_bias) input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups kernels = (kernel_h, kernel_w) strides = (stride_h, stride_w) paddings = (pad_h, pad_w) dilations = (dilation, dilation) assume(height + 2 * paddings[0] >= dilations[0] * (kernels[0] - 1) + 1) assume(width + 2 * paddings[1] >= dilations[1] * (kernels[1] - 1) + 1) input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) weight = torch.rand( (output_channels, input_channels_per_group, kernel_h, kernel_w) ) bias = None if use_bias: bias = torch.rand(output_channels) scripted_conv2d = torch.jit.script( Conv2D(weight, bias, strides, paddings, dilations, groups) ) scripted_conv2d_clamp_prepacked = torch.jit.script( Conv2DPrePacked(weight, bias, strides, paddings, dilations, groups) ) ref_result = scripted_conv2d(input_data) xnnpack_result = scripted_conv2d_clamp_prepacked(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) # Serialize the modules and then deserialize input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) buffer = io.BytesIO() torch.jit.save(scripted_conv2d, buffer) buffer.seek(0) deserialized_conv2d = torch.jit.load(buffer) buffer = io.BytesIO() torch.jit.save(scripted_conv2d_clamp_prepacked, buffer) buffer.seek(0) deserialized_conv2d_clamp_prepacked = torch.jit.load(buffer) ref_result = deserialized_conv2d(input_data) xnnpack_result = deserialized_conv2d_clamp_prepacked(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) @given( batch_size=st.integers(0, 3), input_channels_per_group=st.integers(1, 32), height=st.integers(5, 64), width=st.integers(5, 64), output_channels_per_group=st.integers(1, 32), groups=st.integers(1, 16), kernel_h=st.integers(1, 7), kernel_w=st.integers(1, 7), stride_h=st.integers(1, 2), stride_w=st.integers(1, 2), pad_h=st.integers(0, 2), pad_w=st.integers(0, 2), output_pad_h=st.integers(0, 2), output_pad_w=st.integers(0, 2), dilation=st.integers(1, 2), use_bias=st.booleans(), format=st.sampled_from( [None, torch.preserve_format, torch.contiguous_format, torch.channels_last] ), ) def test_conv2d_transpose( self, batch_size, input_channels_per_group, height, width, output_channels_per_group, groups, kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, output_pad_h, output_pad_w, dilation, use_bias, format, ): class Conv2DT(torch.nn.Module): def __init__( self, weight, bias, strides, paddings, output_paddings, dilations, groups, ): super().__init__() self.weight = weight self.bias = bias self.strides = strides self.paddings = paddings self.output_paddings = output_paddings self.dilations = dilations self.groups = groups def forward(self, x): return F.conv_transpose2d( x, self.weight, self.bias, self.strides, self.paddings, self.output_paddings, self.groups, self.dilations, ) class Conv2DTPrePacked(torch.nn.Module): def __init__( self, weight, bias, strides, paddings, output_paddings, dilations, groups, ): super().__init__() self.packed_weight_bias = ( torch.ops.prepacked.conv2d_transpose_clamp_prepack( weight, bias, strides, paddings, output_paddings, dilations, groups, ) ) def forward(self, x): return torch.ops.prepacked.conv2d_transpose_clamp_run( x, self.packed_weight_bias ) input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups kernels = (kernel_h, kernel_w) strides = (stride_h, stride_w) paddings = (pad_h, pad_w) output_paddings = (output_pad_h, output_pad_w) dilations = (dilation, dilation) assume(height + 2 * paddings[0] >= dilations[0] * (kernels[0] - 1) + 1) assume(width + 2 * paddings[1] >= dilations[1] * (kernels[1] - 1) + 1) assume((output_pad_h < stride_h) and (output_pad_h < dilation)) assume((output_pad_w < stride_w) and (output_pad_w < dilation)) input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) weight = torch.rand( (input_channels, output_channels_per_group, kernel_h, kernel_w) ) bias = None if use_bias: bias = torch.rand(output_channels) scripted_conv2d = torch.jit.script( Conv2DT(weight, bias, strides, paddings, output_paddings, dilations, groups) ) scripted_conv2d_clamp_prepacked = torch.jit.script( Conv2DTPrePacked( weight, bias, strides, paddings, output_paddings, dilations, groups ) ) ref_result = scripted_conv2d(input_data) xnnpack_result = scripted_conv2d_clamp_prepacked(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) # Serialize the modules and then deserialize input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) buffer = io.BytesIO() torch.jit.save(scripted_conv2d, buffer) buffer.seek(0) deserialized_conv2d = torch.jit.load(buffer) buffer = io.BytesIO() torch.jit.save(scripted_conv2d_clamp_prepacked, buffer) buffer.seek(0) deserialized_conv2d_clamp_prepacked = torch.jit.load(buffer) ref_result = deserialized_conv2d(input_data) xnnpack_result = deserialized_conv2d_clamp_prepacked(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) @unittest.skip( "Fails on some platforms, see https://github.com/pytorch/pytorch/issues/73488" ) @given( batch_size=st.integers(0, 3), input_channels_per_group=st.integers(1, 32), height=st.integers(5, 64), width=st.integers(5, 64), output_channels_per_group=st.integers(1, 32), groups=st.integers(1, 16), kernel_h=st.integers(1, 7), kernel_w=st.integers(1, 7), stride_h=st.integers(1, 2), stride_w=st.integers(1, 2), pad_h=st.integers(0, 2), pad_w=st.integers(0, 2), dilation=st.integers(1, 2), linear_weight_output_dim=st.integers(2, 64), use_bias=st.booleans(), format=st.sampled_from( [None, torch.preserve_format, torch.contiguous_format, torch.channels_last] ), ) def test_combined_model( self, batch_size, input_channels_per_group, height, width, output_channels_per_group, groups, kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, dilation, linear_weight_output_dim, use_bias, format, ): class M(torch.nn.Module): def __init__( self, conv_weight, conv_bias, linear_weight, linear_bias, strides, paddings, dilations, groups, ): super().__init__() self.conv_weight = conv_weight self.conv_bias = conv_bias self.linear_weight = linear_weight self.linear_bias = linear_bias self.strides = strides self.paddings = paddings self.dilations = dilations self.groups = groups def forward(self, x): o = F.conv2d( x, self.conv_weight, self.conv_bias, self.strides, self.paddings, self.dilations, self.groups, ) o = o.permute([0, 2, 3, 1]) o = F.linear(o, self.linear_weight, self.linear_bias) return F.relu(o) class MPrePacked(torch.nn.Module): def __init__( self, conv_weight, conv_bias, linear_weight, linear_bias, strides, paddings, dilations, groups, ): super().__init__() self.conv2d_clamp_run_weight_bias = ( torch.ops.prepacked.conv2d_clamp_prepack( conv_weight, conv_bias, strides, paddings, dilations, groups ) ) self.linear_clamp_run_weight_bias = ( torch.ops.prepacked.linear_clamp_prepack(linear_weight, linear_bias) ) def forward(self, x): o = torch.ops.prepacked.conv2d_clamp_run( x, self.conv2d_clamp_run_weight_bias ) o = o.permute([0, 2, 3, 1]) o = torch.ops.prepacked.linear_clamp_run( o, self.linear_clamp_run_weight_bias ) return F.relu(o) input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups kernels = (kernel_h, kernel_w) strides = (stride_h, stride_w) paddings = (pad_h, pad_w) dilations = (dilation, dilation) assume(height + 2 * paddings[0] >= dilations[0] * (kernels[0] - 1) + 1) assume(width + 2 * paddings[1] >= dilations[1] * (kernels[1] - 1) + 1) input_data = torch.rand((batch_size, input_channels, height, width)) if format is not None: input_data = input_data.contiguous(memory_format=format) conv_weight = torch.rand( (output_channels, input_channels_per_group, kernel_h, kernel_w) ) conv_bias = None if use_bias: conv_bias = torch.rand(output_channels) # This is done just to find the output shape of the result # so that the shape of weight for the following linear layer # can be determined. result = F.conv2d( input_data, conv_weight, conv_bias, strides, paddings, dilations, groups ) linear_input_shape = result.shape[1] linear_weight = torch.rand((linear_weight_output_dim, linear_input_shape)) linear_bias = None if use_bias: linear_bias = torch.rand(linear_weight_output_dim) scripted_m = torch.jit.script( M( conv_weight, conv_bias, linear_weight, linear_bias, strides, paddings, dilations, groups, ) ) scripted_m_prepacked = torch.jit.script( MPrePacked( conv_weight, conv_bias, linear_weight, linear_bias, strides, paddings, dilations, groups, ) ) ref_result = scripted_m(input_data) xnnpack_result = scripted_m_prepacked(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) # Serialize the modules and then deserialize input_data = torch.rand((batch_size, input_channels, height, width)) input_data = input_data.contiguous(memory_format=torch.channels_last) buffer = io.BytesIO() torch.jit.save(scripted_m, buffer) buffer.seek(0) deserialized_m = torch.jit.load(buffer) buffer = io.BytesIO() torch.jit.save(scripted_m_prepacked, buffer) buffer.seek(0) deserialized_m_prepacked = torch.jit.load(buffer) ref_result = deserialized_m(input_data) xnnpack_result = deserialized_m_prepacked(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) @unittest.skipUnless( torch.backends.xnnpack.enabled, " XNNPACK must be enabled for these tests." " Please build with USE_XNNPACK=1.", ) @unittest.skipIf( TEST_WITH_TSAN, "TSAN fails with XNNPACK. Does not seem to have a good reason for failures.", ) class TestXNNPACKRewritePass(TestCase): @staticmethod def validate_transformed_module( # To please flake self, pattern_count_map, data_shape, prepack_removal=False, fuse_clamping_ops=False, ): input_data = torch.normal(1, 20, size=data_shape) for jit_method in ["script", "trace"]: module_instance = self if jit_method == "script": scripted_model = torch.jit.script(module_instance) else: scripted_model = torch.jit.trace(module_instance, input_data) scripted_model.eval() ref_result = scripted_model(input_data) torch._C._jit_pass_insert_prepacked_ops(scripted_model._c) if fuse_clamping_ops or prepack_removal: scripted_model._c = torch._C._freeze_module(scripted_model._c) if fuse_clamping_ops: torch._C._jit_pass_fuse_clamp_w_prepacked_linear_conv(scripted_model._c) if prepack_removal: torch._C._jit_pass_fold_prepacking_ops(scripted_model._c) buffer = io.BytesIO() torch.jit.save(scripted_model, buffer) buffer.seek(0) deserialized_scripted_model = torch.jit.load(buffer) for pattern, v in pattern_count_map.items(): if v == 0: FileCheck().check(pattern).run(deserialized_scripted_model.graph) elif v == -1: FileCheck().check_not(pattern).run( deserialized_scripted_model.graph ) else: FileCheck().check_count(pattern, v, exactly=True).run( deserialized_scripted_model.graph ) xnnpack_result = deserialized_scripted_model(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) def test_linear(self): data_shape = [2, 3, 32] weight_output_dim = 24 weight_shape = (weight_output_dim, data_shape[-1]) class Linear(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) def forward(self, x): return F.linear(x, self.weight, self.bias) class LinearNoBias(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(weight_shape), requires_grad=False ) def forward(self, x): return F.linear(x, self.weight, None) # Linear with bias pattern. pattern_count_map = { "Tensor = prim::CallFunction": -1, "prepacked::linear_clamp_prepack": 1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( Linear(), pattern_count_map, data_shape ) TestXNNPACKRewritePass.validate_transformed_module( LinearNoBias(), pattern_count_map, data_shape ) # Conv params batch_size = 2 input_channels_per_group = 6 height = 16 width = 16 output_channels_per_group = 6 groups = 4 kernel_h = kernel_w = 3 stride_h = stride_w = 1 pad_h = pad_w = 1 output_pad_h = output_pad_w = 0 dilation = 1 input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups kernels = (kernel_h, kernel_w) strides = (stride_h, stride_w) paddings = (pad_h, pad_w) output_paddings = (output_pad_h, output_pad_w) dilations = (dilation, dilation) conv_weight_shape = ( output_channels, input_channels_per_group, kernel_h, kernel_w, ) conv_transpose_weight_shape = ( input_channels, output_channels_per_group, kernel_h, kernel_w, ) conv_bias_shape = output_channels class Conv2D(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(conv_weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(conv_bias_shape), requires_grad=False ) self.strides = strides self.paddings = paddings self.dilations = dilations self.groups = groups def forward(self, x): return F.conv2d( x, self.weight, self.bias, self.strides, self.paddings, self.dilations, self.groups, ) class Conv2DT(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(conv_transpose_weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(conv_bias_shape), requires_grad=False ) self.strides = strides self.paddings = paddings self.output_paddings = output_paddings self.dilations = dilations self.groups = groups def forward(self, x): return F.conv_transpose2d( x, self.weight, self.bias, self.strides, self.paddings, self.output_paddings, self.groups, self.dilations, ) data_shape = (batch_size, input_channels, height, width) pattern_count_map = { "Tensor = aten::conv2d": -1, "prepacked::conv2d_clamp_prepack": 1, "prepacked::conv2d_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( Conv2D(), pattern_count_map, data_shape ) transpose_data_shape = (batch_size, input_channels, height, width) transpose_pattern_count_map = { "Tensor = aten::conv_transpose2d": -1, "prepacked::conv2d_transpose_clamp_prepack": 1, "prepacked::conv2d_transpose_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( Conv2DT(), transpose_pattern_count_map, data_shape ) input_data = torch.rand((batch_size, input_channels, height, width)) conv_weight = torch.rand( (output_channels, input_channels_per_group, kernel_h, kernel_w) ) conv_bias = torch.rand(output_channels) result = F.conv2d( input_data, conv_weight, conv_bias, strides, paddings, dilations, groups ) linear_input_shape = result.shape[1] linear_weight_shape = (weight_output_dim, linear_input_shape) class M(torch.nn.Module): def __init__(self, activation_fn=F.relu): super().__init__() self.conv_weight = torch.nn.Parameter( torch.rand(conv_weight_shape), requires_grad=False ) self.conv_bias = torch.nn.Parameter( torch.rand(conv_bias_shape), requires_grad=False ) self.linear_weight = torch.nn.Parameter( torch.rand(linear_weight_shape), requires_grad=False ) self.linear_bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) self.strides = strides self.paddings = paddings self.dilations = dilations self.groups = groups self.activation_fn = activation_fn def forward(self, x): o = F.conv2d( x, self.conv_weight, self.conv_bias, self.strides, self.paddings, self.dilations, self.groups, ) o = self.activation_fn(o) o = o.permute([0, 2, 3, 1]) o = F.linear(o, self.linear_weight, self.linear_bias) return self.activation_fn(o) pattern_count_map = { "Tensor = aten::conv2d": -1, "prepacked::conv2d_clamp_prepack": 1, "prepacked::conv2d_clamp_run": 1, "prepacked::linear_clamp_prepack": 1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( M(), pattern_count_map, data_shape ) pattern_count_map["prepacked::conv2d_clamp_prepack"] = -1 pattern_count_map["Tensor = prim::CallFunction"] = -1 pattern_count_map["prepacked::linear_clamp_prepack"] = -1 TestXNNPACKRewritePass.validate_transformed_module( M(), pattern_count_map, data_shape, prepack_removal=True ) # Not inplace relu fusion test. pattern_count_map = { "aten::relu": 2, "prepacked::conv2d_clamp_prepack": -1, "prepacked::conv2d_clamp_run": 1, "prepacked::linear_clamp_prepack": -1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( M(), pattern_count_map, data_shape, prepack_removal=True ) pattern_count_map["prepacked::conv2d_clamp_prepack"] = -1 pattern_count_map["prepacked::linear_clamp_prepack"] = -1 pattern_count_map["aten::relu"] = -1 TestXNNPACKRewritePass.validate_transformed_module( M(), pattern_count_map, data_shape, prepack_removal=True, fuse_clamping_ops=True, ) # Inplace relu fusion test. pattern_count_map = { "aten::relu": 2, "prepacked::conv2d_clamp_prepack": -1, "prepacked::conv2d_clamp_run": 1, "prepacked::linear_clamp_prepack": -1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( M(F.relu_), pattern_count_map, data_shape, prepack_removal=True ) pattern_count_map["prepacked::conv2d_clamp_prepack"] = -1 pattern_count_map["prepacked::linear_clamp_prepack"] = -1 pattern_count_map["aten::relu"] = -1 TestXNNPACKRewritePass.validate_transformed_module( M(F.relu_), pattern_count_map, data_shape, prepack_removal=True, fuse_clamping_ops=True, ) # Not inplace hardtanh fusion test. pattern_count_map = { "aten::hardtanh": 2, "prepacked::conv2d_clamp_prepack": -1, "prepacked::conv2d_clamp_run": 1, "prepacked::linear_clamp_prepack": -1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( M(F.hardtanh), pattern_count_map, data_shape, prepack_removal=True ) pattern_count_map["prepacked::conv2d_clamp_prepack"] = -1 pattern_count_map["prepacked::linear_clamp_prepack"] = -1 pattern_count_map["aten::hardtanh"] = -1 TestXNNPACKRewritePass.validate_transformed_module( M(F.hardtanh), pattern_count_map, data_shape, prepack_removal=True, fuse_clamping_ops=True, ) # Inplace hardtanh fusion test. pattern_count_map = { "aten::hardtanh_": 2, "prepacked::conv2d_clamp_prepack": -1, "prepacked::conv2d_clamp_run": 1, "prepacked::linear_clamp_prepack": -1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( M(F.hardtanh_), pattern_count_map, data_shape, prepack_removal=True ) pattern_count_map["prepacked::conv2d_clamp_prepack"] = -1 pattern_count_map["prepacked::linear_clamp_prepack"] = -1 pattern_count_map["aten::hardtanh_"] = -1 TestXNNPACKRewritePass.validate_transformed_module( M(F.hardtanh_), pattern_count_map, data_shape, prepack_removal=True, fuse_clamping_ops=True, ) class MFusionAntiPattern(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear_weight = torch.nn.Parameter( torch.rand(linear_weight_shape), requires_grad=False ) self.linear_bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) self.strides = strides self.paddings = paddings self.dilations = dilations self.groups = groups def forward(self, x): o = F.linear(x, self.linear_weight, self.linear_bias) o = F.relu(o) o = F.hardtanh(o) return o # Unfusable hardtanh. pattern_count_map = { "aten::hardtanh": 1, # hardtanh cannot be. "aten::relu": -1, # relu is fused. "prepacked::linear_clamp_prepack": -1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( MFusionAntiPattern(), pattern_count_map, (16, linear_weight_shape[1]), prepack_removal=True, fuse_clamping_ops=True, ) class MFusionAntiPatternParamMinMax(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear_weight = torch.nn.Parameter( torch.rand(linear_weight_shape), requires_grad=False ) self.linear_bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) self.strides = strides self.paddings = paddings self.dilations = dilations self.groups = groups def forward(self, x): min = x[0, 0] max = min + 10 o = F.linear(x, self.linear_weight, self.linear_bias) o = F.hardtanh(o, min, max) return o # Unfusable hardtanh. pattern_count_map = { "aten::hardtanh": 1, # hardtanh cannot be. "prepacked::linear_clamp_prepack": -1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( MFusionAntiPatternParamMinMax(), pattern_count_map, (16, linear_weight_shape[1]), prepack_removal=True, fuse_clamping_ops=True, ) def test_decomposed_linear(self): data_shape = [2, 32] weight_output_dim = 24 weight_shape = (weight_output_dim, data_shape[-1]) class DecomposedLinearAddmm(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) def forward(self, x): weight_t = self.weight.t() return torch.addmm(self.bias, x, weight_t) class DecomposedLinearMatmulAdd(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) def forward(self, x): weight_t = self.weight.t() y = torch.matmul(x, weight_t) res = y.add_(self.bias) return res class DecomposedLinearMatmul(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(weight_output_dim), requires_grad=False ) def forward(self, x): weight_t = self.weight.t() res = torch.matmul(x, weight_t) return res # Linear with bias pattern. pattern_count_map = { "Tensor = prim::CallFunction": -1, "prepacked::linear_clamp_prepack": 1, "prepacked::linear_clamp_run": 1, } TestXNNPACKRewritePass.validate_transformed_module( DecomposedLinearAddmm(), pattern_count_map, data_shape ) TestXNNPACKRewritePass.validate_transformed_module( DecomposedLinearMatmulAdd(), pattern_count_map, data_shape ) TestXNNPACKRewritePass.validate_transformed_module( DecomposedLinearMatmul(), pattern_count_map, data_shape ) @unittest.skipUnless( torch.backends.xnnpack.enabled, " XNNPACK must be enabled for these tests." " Please build with USE_XNNPACK=1.", ) @unittest.skipIf( TEST_WITH_TSAN, "TSAN is not fork-safe since we're forking in a multi-threaded environment", ) class TestXNNPACKConv1dTransformPass(TestCase): @staticmethod def validate_transform_conv1d_to_conv2d( self, pattern_count_transformed_map, pattern_count_optimized_map, data_shape ): input_data = torch.normal(1, 20, size=data_shape) for jit_method in ["script", "trace"]: module_instance = self if jit_method == "script": scripted_model = torch.jit.script(module_instance) else: scripted_model = torch.jit.trace(module_instance, input_data) scripted_model.eval() ref_result = scripted_model(input_data) torch._C._jit_pass_transform_conv1d_to_conv2d(scripted_model._c) optimized_scripted_model = optimize_for_mobile(scripted_model) buffer = io.BytesIO() torch.jit.save(scripted_model, buffer) buffer.seek(0) deserialized_scripted_model = torch.jit.load(buffer) for pattern, v in pattern_count_transformed_map.items(): if v == 0: FileCheck().check(pattern).run(deserialized_scripted_model.graph) elif v == -1: FileCheck().check_not(pattern).run( deserialized_scripted_model.graph ) else: FileCheck().check_count(pattern, v, exactly=True).run( deserialized_scripted_model.graph ) transformed_result = deserialized_scripted_model(input_data) torch.testing.assert_close( ref_result, transformed_result, rtol=1e-2, atol=1e-3 ) optimized_buffer = io.BytesIO() torch.jit.save(optimized_scripted_model, optimized_buffer) optimized_buffer.seek(0) deserialized_optimized_scripted_model = torch.jit.load(optimized_buffer) for pattern, v in pattern_count_optimized_map.items(): if v == 0: FileCheck().check(pattern).run( deserialized_optimized_scripted_model.graph ) elif v == -1: FileCheck().check_not(pattern).run( deserialized_optimized_scripted_model.graph ) else: FileCheck().check_count(pattern, v, exactly=True).run( deserialized_optimized_scripted_model.graph ) xnnpack_result = deserialized_optimized_scripted_model(input_data) torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3) @unittest.skipIf(IS_FBCODE, "T137513244") def test_conv1d_basic(self): batch_size_list = range(1, 3) input_channels_per_group_list = range(10, 12) width_list = range(10, 12) output_channels_per_group_list = range(10, 12) groups_list = range(1, 3) kernel_list = range(1, 4) stride_list = range(1, 3) padding_list = range(0, 3) dilation_list = range(1, 3) for hparams in itertools.product( batch_size_list, input_channels_per_group_list, width_list, output_channels_per_group_list, groups_list, kernel_list, stride_list, padding_list, dilation_list, ): ( batch_size, input_channels_per_group, width, output_channels_per_group, groups, kernel, stride, padding, dilation, ) = hparams input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups conv_weight_shape = (output_channels, input_channels_per_group, kernel) conv_bias_shape = output_channels class Conv1D(torch.nn.Module): def __init__(self) -> None: super().__init__() self.weight = torch.nn.Parameter( torch.rand(conv_weight_shape), requires_grad=False ) self.bias = torch.nn.Parameter( torch.rand(conv_bias_shape), requires_grad=False ) self.stride = stride self.padding = padding self.dilation = dilation self.groups = groups def forward(self, x): return F.conv1d( x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups, ) data_shape = (batch_size, input_channels, width) pattern_count_transformed_map = { "Tensor = aten::conv1d": -1, "Tensor = aten::conv2d": 1, } pattern_count_optimized_map = { "Tensor = aten::conv1d": -1, "Tensor = aten::conv2d": -1, "prepacked::conv2d_clamp_prepack": -1, "prepacked::conv2d_clamp_run": 1, } TestXNNPACKConv1dTransformPass.validate_transform_conv1d_to_conv2d( Conv1D(), pattern_count_transformed_map, pattern_count_optimized_map, data_shape, ) # See https://github.com/pytorch/pytorch/issues/46066 @slowTest def test_conv1d_with_relu_fc(self): batch_size_list = range(1, 3) input_channels_per_group_list = range(10, 12) width_list = range(10, 12) output_channels_per_group_list = range(10, 12) groups_list = range(1, 3) kernel_list = range(1, 4) stride_list = range(1, 3) padding_list = range(0, 3) dilation_list = range(1, 3) output_features_list = range(1, 3) for hparams in itertools.product( batch_size_list, input_channels_per_group_list, width_list, output_channels_per_group_list, groups_list, kernel_list, stride_list, padding_list, dilation_list, output_features_list, ): ( batch_size, input_channels_per_group, width, output_channels_per_group, groups, kernel, stride, padding, dilation, output_features, ) = hparams input_channels = input_channels_per_group * groups output_channels = output_channels_per_group * groups conv_weight_shape = (output_channels, input_channels_per_group, kernel) conv_bias_shape = output_channels conv_output_width = ( int((width + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1 ) fc_weight_shape = (output_features, output_channels * conv_output_width) fc_bias_shape = output_features class Net(torch.nn.Module): def __init__(self) -> None: super().__init__() self.conv_weight = torch.nn.Parameter( torch.rand(conv_weight_shape), requires_grad=False ) self.conv_bias = torch.nn.Parameter( torch.rand(conv_bias_shape), requires_grad=False ) self.stride = stride self.padding = padding self.dilation = dilation self.groups = groups self.fc_weight = torch.nn.Parameter( torch.rand(fc_weight_shape), requires_grad=False ) self.fc_bias = torch.nn.Parameter( torch.rand(fc_bias_shape), requires_grad=False ) def forward(self, x): x = F.conv1d( x, self.conv_weight, self.conv_bias, self.stride, self.padding, self.dilation, self.groups, ) x = F.relu(x) x = x.view(x.size(0), -1) x = F.linear(x, self.fc_weight, self.fc_bias) return x data_shape = (batch_size, input_channels, width) pattern_count_transformed_map = { "Tensor = aten::conv1d": -1, "Tensor = aten::conv2d": 1, } pattern_count_optimized_map = { "Tensor = aten::conv1d": -1, "Tensor = aten::conv2d": -1, "prepacked::conv2d_clamp_prepack": -1, "prepacked::conv2d_clamp_run": 1, } TestXNNPACKConv1dTransformPass.validate_transform_conv1d_to_conv2d( Net(), pattern_count_transformed_map, pattern_count_optimized_map, data_shape, ) if __name__ == "__main__": run_tests()