| # Owner(s): ["module: fx"] |
| |
| import copy |
| import sys |
| import logging |
| from typing import List, Tuple |
| |
| import torch |
| from torch.fx._symbolic_trace import symbolic_trace |
| from torch.fx.experimental.proxy_tensor import make_fx |
| from torch.fx.passes.backends.nvfuser import NvFuserBackend |
| |
| from torch.testing._internal.common_utils import run_tests, TEST_CUDA, TestCase |
| from torch.testing._internal.common_device_type import ( |
| instantiate_device_type_tests, |
| dtypes, |
| ) |
| |
| if not TEST_CUDA: |
| print('CUDA not available, skipping tests', file=sys.stderr) |
| TestCase = object # noqa: F811 |
| |
| logging.basicConfig(level=logging.DEBUG) |
| logger = logging.getLogger(__name__) |
| |
| class HF_T5_Partial(torch.nn.Module): |
| |
| def inputs_meta(self): |
| return [ |
| (torch.Size([512, 512]), torch.float32), |
| (torch.Size([512, 512]), torch.float32), |
| (torch.Size([512, 512]), torch.float32), |
| (torch.Size([512, 512]), torch.float32), |
| (torch.Size([512]), torch.float32), |
| (torch.Size([2048, 512]), torch.float32), |
| (torch.Size([512, 2048]), torch.float32), |
| (torch.Size([512]), torch.float32), |
| (torch.Size([8, 1024, 512]), torch.float32), |
| (torch.Size([8, 8, 1024, 1024]), torch.float32), |
| ] |
| |
| def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, |
| primals_6, primals_7, primals_8, primals_9, primals_10): |
| pow_1 = torch.ops.aten.pow(primals_9, 2) |
| mean = torch.ops.aten.mean(pow_1, [-1], True) |
| add = torch.ops.aten.add(mean, 1e-06) |
| rsqrt = torch.ops.aten.rsqrt(add) |
| mul = torch.ops.aten.mul(primals_9, rsqrt) |
| mul_1 = torch.ops.aten.mul(primals_5, mul) |
| t = torch.ops.aten.t(primals_3) |
| view = torch.ops.aten.view(mul_1, [8192, 512]) |
| mm = torch.ops.aten.mm(view, t) |
| _unsafe_view = torch.ops.aten._unsafe_view(mm, [8, 1024, 512]) |
| view_1 = torch.ops.aten.view(_unsafe_view, [8, -1, 8, 64]) |
| transpose = torch.ops.aten.transpose(view_1, 1, 2) |
| t_1 = torch.ops.aten.t(primals_1) |
| view_2 = torch.ops.aten.view(mul_1, [8192, 512]) |
| mm_1 = torch.ops.aten.mm(view_2, t_1) |
| _unsafe_view_1 = torch.ops.aten._unsafe_view(mm_1, [8, 1024, 512]) |
| view_3 = torch.ops.aten.view(_unsafe_view_1, [8, -1, 8, 64]) |
| transpose_1 = torch.ops.aten.transpose(view_3, 1, 2) |
| t_2 = torch.ops.aten.t(primals_4) |
| view_4 = torch.ops.aten.view(mul_1, [8192, 512]) |
| mm_2 = torch.ops.aten.mm(view_4, t_2) |
| _unsafe_view_2 = torch.ops.aten._unsafe_view(mm_2, [8, 1024, 512]) |
| view_5 = torch.ops.aten.view(_unsafe_view_2, [8, -1, 8, 64]) |
| transpose_2 = torch.ops.aten.transpose(view_5, 1, 2) |
| transpose_3 = torch.ops.aten.transpose(transpose_1, 3, 2) |
| expand = torch.ops.aten.expand(transpose, [8, 8, 1024, 64]) |
| clone = torch.ops.aten.clone(expand, memory_format=torch.contiguous_format) |
| _unsafe_view_3 = torch.ops.aten._unsafe_view(clone, [64, 1024, 64]) |
| expand_1 = torch.ops.aten.expand(transpose_3, [8, 8, 64, 1024]) |
| clone_1 = torch.ops.aten.clone(expand_1, memory_format=torch.contiguous_format) |
| _unsafe_view_4 = torch.ops.aten._unsafe_view(clone_1, [64, 64, 1024]) |
| bmm = torch.ops.aten.bmm(_unsafe_view_3, _unsafe_view_4) |
| _unsafe_view_5 = torch.ops.aten._unsafe_view(bmm, [8, 8, 1024, 1024]) |
| add_ = torch.ops.aten.add_(_unsafe_view_5, primals_10) |
| _softmax = torch.ops.aten._softmax(add_, -1, False) |
| expand_2 = torch.ops.aten.expand(_softmax, [8, 8, 1024, 1024]) |
| view_6 = torch.ops.aten.view(expand_2, [64, 1024, 1024]) |
| expand_3 = torch.ops.aten.expand(transpose_2, [8, 8, 1024, 64]) |
| clone_2 = torch.ops.aten.clone(expand_3, memory_format=torch.contiguous_format) |
| _unsafe_view_6 = torch.ops.aten._unsafe_view(clone_2, [64, 1024, 64]) |
| bmm_1 = torch.ops.aten.bmm(view_6, _unsafe_view_6) |
| _unsafe_view_7 = torch.ops.aten._unsafe_view(bmm_1, [8, 8, 1024, 64]) |
| transpose_4 = torch.ops.aten.transpose(_unsafe_view_7, 1, 2) |
| clone_3 = torch.ops.aten.clone(transpose_4, memory_format=torch.contiguous_format) |
| view_7 = torch.ops.aten.view(clone_3, [8, -1, 512]) |
| t_3 = torch.ops.aten.t(primals_2) |
| view_8 = torch.ops.aten.view(view_7, [8192, 512]) |
| mm_3 = torch.ops.aten.mm(view_8, t_3) |
| _unsafe_view_8 = torch.ops.aten._unsafe_view(mm_3, [8, 1024, 512]) |
| add_1 = torch.ops.aten.add(primals_9, _unsafe_view_8) |
| pow_2 = torch.ops.aten.pow(add_1, 2) |
| mean_1 = torch.ops.aten.mean(pow_2, [-1], True) |
| add_2 = torch.ops.aten.add(mean_1, 1e-06) |
| rsqrt_1 = torch.ops.aten.rsqrt(add_2) |
| mul_2 = torch.ops.aten.mul(add_1, rsqrt_1) |
| mul_3 = torch.ops.aten.mul(primals_8, mul_2) |
| t_4 = torch.ops.aten.t(primals_6) |
| view_9 = torch.ops.aten.view(mul_3, [8192, 512]) |
| mm_4 = torch.ops.aten.mm(view_9, t_4) |
| _unsafe_view_9 = torch.ops.aten._unsafe_view(mm_4, [8, 1024, 2048]) |
| relu = torch.ops.aten.relu(_unsafe_view_9) |
| t_5 = torch.ops.aten.t(primals_7) |
| view_10 = torch.ops.aten.view(relu, [8192, 2048]) |
| mm_5 = torch.ops.aten.mm(view_10, t_5) |
| _unsafe_view_10 = torch.ops.aten._unsafe_view(mm_5, [8, 1024, 512]) |
| add_3 = torch.ops.aten.add(add_1, _unsafe_view_10) |
| return [add_3, rsqrt, _unsafe_view_3, t_3, _softmax, view_6, mul_2, t, view_9, t_1, primals_5, add_1, |
| _unsafe_view_4, view_2, view_10, t_5, t_2, primals_8, view_4, view_8, rsqrt_1, primals_9, t_4, |
| mul, _unsafe_view_6, relu, view] |
| |
| |
| class TestFxNvFuserBackend(TestCase): |
| |
| def _generate_random_inputs(self, device, inputs_meta: List[Tuple[torch.Size, torch.dtype]]): |
| inputs = [] |
| for meta in inputs_meta: |
| shape, dtype = meta |
| |
| if dtype in {torch.int, torch.int32, torch.int64, torch.bool, torch.int, torch.uint8}: |
| input = torch.randint(0, 1, shape, dtype=dtype, device=device) |
| else: |
| input = torch.rand(shape, dtype=dtype, device=device) |
| |
| inputs.append(input) |
| |
| return inputs |
| |
| |
| @dtypes(torch.float32) |
| def test_nvfuser_call_module_backend(self, device, dtype): |
| |
| class Model(torch.nn.Module): |
| |
| def __init__(self): |
| super(Model, self).__init__() |
| self.bn = torch.nn.BatchNorm2d(3) |
| self.relu = torch.nn.ReLU() |
| |
| def forward(self, inp): |
| o = self.bn(inp) |
| o = self.relu(o) |
| return o |
| |
| inp = torch.randn(2, 3, 4, 5).to(dtype=dtype, device=device) |
| m = Model().to(dtype=dtype, device=device) |
| |
| # note that the traced module here contains only `call_module` node, |
| # which isn't fused by nvfuser backend. But `nvfuser.compile` should run without error |
| traced = symbolic_trace(m) |
| |
| nvfuser = NvFuserBackend() |
| compiled_module = nvfuser.compile(traced) |
| |
| eager_result = m(inp) |
| nvfuser_result = compiled_module(inp) |
| |
| torch.testing.assert_close(eager_result, nvfuser_result, rtol=1e-5, atol=1e-5) |
| |
| |
| @dtypes(torch.float32) |
| def test_nvfuser_backend(self, device, dtype): |
| m = HF_T5_Partial() |
| m.to(device) |
| |
| traced = symbolic_trace(m) |
| |
| nvfuser = NvFuserBackend() |
| compiled_module = nvfuser.compile(traced) |
| |
| inputs = self._generate_random_inputs(device, m.inputs_meta()) |
| |
| eager_result = m(*inputs) |
| nvfuser_result = compiled_module(*inputs) |
| |
| torch.testing.assert_close(eager_result, nvfuser_result, rtol=1e-5, atol=1e-5) |
| |
| @dtypes(torch.float32) |
| def test_aten_square(self, device, dtype): |
| |
| def fn(x): |
| square = torch.square(x) |
| a = square + 1 |
| b = a + 1 |
| return b |
| |
| inputs = torch.randn(4, device=device) |
| traced = make_fx(fn)(inputs) |
| |
| nvfuser = NvFuserBackend() |
| compiled_module = nvfuser.compile(copy.deepcopy(traced)) |
| |
| for node in compiled_module.graph.nodes: |
| if node.op == "call_function": |
| assert "fused" in str(node.target), "the entire function should be fused into a single fusion group" |
| |
| eager_result = traced(inputs) |
| nvfuser_result = compiled_module(inputs) |
| torch.testing.assert_close(eager_result, nvfuser_result, rtol=1e-5, atol=1e-5) |
| |
| @dtypes(torch.float32) |
| def test_aten_leakyrelu(self, device, dtype): |
| |
| def fn(x): |
| square = torch.ops.aten.leaky_relu(x, 0.1) |
| a = square + 1 |
| b = a + 1 |
| return b |
| |
| inputs = torch.randn(4, device=device) |
| traced = make_fx(fn)(inputs) |
| |
| nvfuser = NvFuserBackend() |
| compiled_module = nvfuser.compile(copy.deepcopy(traced)) |
| |
| for node in compiled_module.graph.nodes: |
| if node.op == "call_function": |
| assert "fused" in str(node.target), "the entire function should be fused into a single fusion group" |
| |
| eager_result = traced(inputs) |
| nvfuser_result = compiled_module(inputs) |
| torch.testing.assert_close(eager_result, nvfuser_result, rtol=1e-5, atol=1e-5) |
| |
| @dtypes(torch.float32) |
| def test_aten_where(self, device, dtype): |
| |
| def fn(x): |
| where = torch.ops.aten.where(x < 0, -x, x) |
| a = where + 1 |
| b = a + 1 |
| return b |
| |
| inputs = torch.randn(4, device=device) |
| traced = make_fx(fn)(inputs) |
| |
| nvfuser = NvFuserBackend() |
| compiled_module = nvfuser.compile(copy.deepcopy(traced)) |
| |
| for node in compiled_module.graph.nodes: |
| if node.op == "call_function": |
| assert "fused" in str(node.target), "the entire function should be fused into a single fusion group" |
| |
| eager_result = traced(inputs) |
| nvfuser_result = compiled_module(inputs) |
| torch.testing.assert_close(eager_result, nvfuser_result, rtol=1e-5, atol=1e-5) |
| |
| instantiate_device_type_tests(TestFxNvFuserBackend, globals(), only_for="cuda") |
| |
| if __name__ == "__main__": |
| run_tests() |
