| # Owner(s): ["module: nvfuser"] |
| |
| import unittest |
| from typing import List |
| |
| import torch |
| from torch.testing._internal.common_utils import run_tests, TEST_WITH_ROCM, TestCase |
| from torch.testing._internal.jit_utils import RUN_CUDA |
| import torch._refs as refs |
| import torch._prims as prims |
| |
| # Will only create the _nvfuser module if CUDA is available |
| if hasattr(torch._C, "_nvfuser"): |
| from torch._C._nvfuser import Fusion, FusionCache, FusionDefinition, DataType |
| |
| RUN_NVFUSER = RUN_CUDA and not TEST_WITH_ROCM |
| |
| def is_pre_volta(): |
| if not RUN_NVFUSER: |
| return False |
| prop = torch.cuda.get_device_properties(torch.cuda.current_device()) |
| return prop.major < 7 |
| |
| @unittest.skipIf(not RUN_NVFUSER, "requires CUDA") |
| @unittest.skipIf(is_pre_volta(), "Only supported on Volta and newer devices.") |
| class TestNvFuserFrontend(TestCase): |
| def test_basic(self) : |
| input1 = torch.ones(2, 4, 8, device='cuda') |
| input2 = torch.ones(2, 4, 8, device='cuda') |
| fc = FusionCache.get() |
| before_fusions = fc.num_fusions() |
| |
| fs1 = Fusion() |
| with FusionDefinition(fs1) as fd : |
| t0 = fd.define_tensor(3) |
| t1 = fd.define_tensor(3) |
| c0 = fd.define_constant(3.0) |
| |
| t2 = fd.ops.add(t0, t1) |
| t3 = fd.ops.mul(t2, c0) |
| t4 = fd.ops.sum(t3, [-1], False, DataType.Float) |
| |
| fd.add_output(t4) |
| |
| # Expected Output is a tensor of 48's |
| nvf_out1 = fs1.execute([input1, input2])[0] |
| |
| # Create a new fusion with the same definition, it should hit the cache! |
| fs2 = Fusion() |
| with FusionDefinition(fs2) as fd : |
| t0 = fd.define_tensor(3) |
| t1 = fd.define_tensor(3) |
| c0 = fd.define_constant(3.0) |
| |
| t2 = fd.ops.add(t0, t1) |
| t3 = fd.ops.mul(t2, c0) |
| t4 = fd.ops.sum(t3, [-1], False, DataType.Float) |
| |
| fd.add_output(t4) |
| |
| nvf_out2 = fs2.execute([input1, input2])[0] |
| |
| # Check there is still only 1 cache entry |
| fc = FusionCache.get() |
| self.assertEqual(fc.num_fusions() - before_fusions, 1) |
| |
| # Create a fusion from a fusion id and make sure it executes! |
| fs3 = Fusion(fs2.id()) |
| nvf_out3 = fs3.execute([input1, input2])[0] |
| |
| eager_out = torch.sum((input1 + input2) * 3.0, dim=-1) |
| self.assertEqual(eager_out, nvf_out1) |
| self.assertEqual(eager_out, nvf_out2) |
| self.assertEqual(eager_out, nvf_out3) |
| |
| def test_basic_fp16(self) : |
| fs = Fusion() |
| with FusionDefinition(fs) as fd : |
| t0 = fd.define_tensor(3, DataType.Half) |
| t1 = fd.define_tensor(3, DataType.Half) |
| c0 = fd.define_constant(3.0) |
| |
| t2 = fd.ops.add(t0, t1) |
| t3 = fd.ops.mul(t2, c0) |
| t4 = fd.ops.sum(t3, [-1], False, DataType.Float) |
| |
| t5 = fd.ops.cast(t4, DataType.Half) |
| fd.add_output(t5) |
| |
| input1 = torch.ones(2, 4, 8, device='cuda', dtype=torch.float16) |
| input2 = torch.ones(2, 4, 8, device='cuda', dtype=torch.float16) |
| |
| # Expected Output is a tensor of 48's |
| nvf_out = fs.execute([input1, input2])[0] |
| eager_out = torch.sum((input1 + input2) * 3.0, dim=-1) |
| self.assertEqual(eager_out, nvf_out) |
| |
| def test_cast_double_to_half(self) : |
| fs = Fusion() |
| with FusionDefinition(fs) as fd : |
| t0 = fd.define_tensor(2, DataType.Double) |
| t1 = fd.define_tensor(2, DataType.Double) |
| |
| t0h = fd.ops.cast(t0, DataType.Half) |
| t1h = fd.ops.cast(t1, DataType.Half) |
| t2 = fd.ops.add(t0h, t1h) |
| t3 = fd.ops.relu(t2) |
| t4 = fd.ops.cast(t3, DataType.Half) |
| |
| fd.add_output(t4) |
| |
| input1 = torch.randn(2, 4, device='cuda', dtype=torch.float64) |
| input2 = torch.randn(2, 4, device='cuda', dtype=torch.float64) |
| |
| nvf_out = fs.execute([input1, input2])[0] |
| eager_out = torch.relu(input1.to(torch.half) + input2.to(torch.half)) |
| self.assertEqual(eager_out, nvf_out) |
| |
| def test_promote_to_double(self) : |
| fs = Fusion() |
| |
| with FusionDefinition(fs) as fd : |
| t0 = fd.define_tensor(2, DataType.Half) |
| t1 = fd.define_tensor(2, DataType.Double) |
| |
| t2 = fd.ops.add(t0, t1) |
| t5 = fd.ops.relu(t2) |
| |
| fd.add_output(t5) |
| |
| input1 = torch.randn(2, 4, device='cuda', dtype=torch.float16) |
| input2 = torch.randn(2, 4, device='cuda', dtype=torch.float64) |
| |
| nvf_out = fs.execute([input1, input2])[0] |
| eager_out = torch.relu(input1 + input2) |
| self.assertEqual(eager_out, nvf_out) |
| |
| def test_implicit_broadcast_input(self) : |
| fs = Fusion() |
| with FusionDefinition(fs) as fd : |
| t0 = fd.define_tensor(1) |
| t1 = fd.define_tensor(3) |
| |
| t0_b = fd.ops.broadcast_in_dim(t0, [2, 3, 4], [1]) |
| t2 = fd.ops.add(t0_b, t1) |
| |
| fd.add_output(t2) |
| |
| input1 = torch.randn(3, device='cuda') |
| input2 = torch.randn(2, 3, 4, device='cuda') |
| |
| nvf_out = fs.execute([input1, input2])[0] |
| eager_out = refs.add(prims.broadcast_in_dim(input1, [2, 3, 4], [1]), input2) |
| self.assertEqual(eager_out, nvf_out) |
| |
| def test_explicit_broadcast_input(self) : |
| input1 = torch.randn(1, 1, 4, device='cuda') |
| input2 = torch.randn(2, 3, 4, device='cuda') |
| |
| fs = Fusion() |
| with FusionDefinition(fs) as fd : |
| t0 = fd.define_tensor(sizes=input1.size(), strides=input1.stride()) |
| t1 = fd.define_tensor(sizes=input2.size(), strides=input2.stride()) |
| |
| t0_b = fd.ops.broadcast_in_dim(t0, [2, 3, 4], [0, 1, 2]) |
| t2 = fd.ops.add(t0_b, t1) |
| |
| fd.add_output(t2) |
| |
| nvf_out = fs.execute([input1, input2])[0] |
| eager_out = refs.add(prims.broadcast_in_dim(input1, [2, 3, 4], [0, 1, 2]), input2) |
| self.assertEqual(eager_out, nvf_out) |
| |
| def test_broadcast_mixing(self) : |
| fs = Fusion() |
| with FusionDefinition(fs) as fd : |
| t0 = fd.define_tensor([3, 1], [1, 1]) |
| t1 = fd.define_tensor(1) |
| |
| t1_b = fd.ops.broadcast_in_dim(t1, [3, 3], [0]) |
| t2 = fd.ops.add(t0, t1_b) |
| |
| fd.add_output(t2) |
| |
| input1 = torch.randn(3, 1, device='cuda') |
| input2 = torch.randn(3, device='cuda') |
| |
| nvf_out = fs.execute([input1, input2])[0] |
| eager_out = refs.add(input1, prims.broadcast_in_dim(input2, [3, 3], [0])) |
| self.assertEqual(eager_out, nvf_out) |
| |
| def test_prim_layer_norm_fwd(self) : |
| def primitive_definition( |
| inputs: torch.Tensor, |
| weight: torch.Tensor, |
| bias: torch.Tensor, |
| normalization_axis: int, |
| keepdim: bool, |
| ) -> torch.Tensor: |
| mean = inputs.mean(normalization_axis, keepdim=keepdim) |
| diff = inputs - mean |
| diff_sq = diff * diff |
| var = diff_sq.mean(normalization_axis, keepdim=keepdim) |
| pre_shift_scale_norm_output = (inputs - mean) / torch.sqrt(var + 1e-12) |
| norm_output = weight * pre_shift_scale_norm_output + bias |
| return norm_output |
| |
| def nvfuser_fusion( |
| fd: FusionDefinition, |
| normalization_axis: int, |
| norm_size: int, |
| input_shape: List[int], |
| eps: float, |
| keepDim: bool |
| ) -> None : |
| inputs = fd.define_tensor(symbolic_sizes=[-1, -1, -1], contiguous=[True, True, True], dtype=DataType.Float) |
| weights = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float) |
| bias = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float) |
| sum0 = fd.ops.sum(inputs, axes=[normalization_axis], keepdim=keepDim) |
| norm_const = fd.define_constant(norm_size) |
| mean = fd.ops.div(sum0, norm_const) |
| diff = fd.ops.sub(inputs, mean) |
| diff_sq = fd.ops.mul(diff, diff) |
| sum1 = fd.ops.sum(diff_sq, axes=[normalization_axis], keepdim=keepDim) |
| var = fd.ops.div(sum1, norm_const) |
| eps_const = fd.define_constant(eps) |
| var_eps = fd.ops.add(var, eps_const) |
| invstd = fd.ops.rsqrt(var_eps) |
| pre_scale_bias = fd.ops.mul(diff, invstd) |
| weights_bcast = fd.ops.broadcast_in_dim(weights, output_shape=input_shape, broadcast_dims=[2]) |
| scale = fd.ops.mul(pre_scale_bias, weights_bcast) |
| bias_bcast = fd.ops.broadcast_in_dim(bias, output_shape=input_shape, broadcast_dims=[2]) |
| out = fd.ops.add(scale, bias_bcast) |
| fd.add_output(out) |
| fd.add_output(mean) |
| fd.add_output(invstd) |
| |
| def nvfuser_fusion_var_mean( |
| fd: FusionDefinition, |
| normalization_axis: int, |
| norm_size: int, |
| input_shape: List[int], |
| eps: float, |
| keepDim: bool |
| ) -> None : |
| inputs = fd.define_tensor(symbolic_sizes=[-1, -1, -1], contiguous=[True, True, True], dtype=DataType.Float) |
| weights = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float) |
| bias = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float) |
| var, mean = fd.ops.var_mean(inputs, axes=[normalization_axis], correction=0, keepdim=keepDim) |
| eps_const = fd.define_constant(eps) |
| var_eps = fd.ops.add(var, eps_const) |
| invstd = fd.ops.rsqrt(var_eps) |
| diff = fd.ops.sub(inputs, mean) |
| pre_scale_bias = fd.ops.mul(diff, invstd) |
| weights_bcast = fd.ops.broadcast_in_dim(weights, output_shape=input_shape, broadcast_dims=[2]) |
| scale = fd.ops.mul(pre_scale_bias, weights_bcast) |
| bias_bcast = fd.ops.broadcast_in_dim(bias, output_shape=input_shape, broadcast_dims=[2]) |
| out = fd.ops.add(scale, bias_bcast) |
| fd.add_output(out) |
| fd.add_output(mean) |
| fd.add_output(invstd) |
| |
| input_size = [64, 128, 1024] |
| dtype = torch.float32 |
| device = 'cuda' |
| inputs = torch.randn(*input_size, device=device, requires_grad=True) |
| weights = torch.nn.Parameter(torch.randn(input_size[2], dtype=dtype, device=device)) |
| biases = torch.nn.Parameter(torch.randn(input_size[2], dtype=dtype, device=device)) |
| fc = FusionCache.get() |
| before_fusions = fc.num_fusions() |
| |
| for _ in range(5) : |
| nvf_fusion = Fusion() |
| with FusionDefinition(nvf_fusion) as fd: |
| nvfuser_fusion(fd, 2, inputs.size()[2], inputs.size(), 1e-12, True) |
| nvf_out = nvf_fusion.execute([inputs, weights, biases]) |
| |
| for _ in range(5) : |
| nvf_var_mean_fusion = Fusion() |
| with FusionDefinition(nvf_var_mean_fusion) as fd: |
| nvfuser_fusion_var_mean(fd, 2, inputs.size()[2], inputs.size(), 1e-12, True) |
| nvf_var_mean_out = nvf_var_mean_fusion.execute([inputs, weights, biases]) |
| |
| for _ in range(5) : |
| eager_out = primitive_definition(inputs, weights, biases, 2, True) |
| |
| self.assertEqual(eager_out, nvf_out[0]) |
| self.assertEqual(eager_out, nvf_var_mean_out[0]) |
| fusion_cache = FusionCache.get() |
| self.assertEqual(fc.num_fusions() - before_fusions, 2) |
| |
| def test_prim_rms_norm_fwd(self) : |
| def primitive_definition( |
| inputs: torch.Tensor, |
| weight: torch.Tensor, |
| normalization_axis: int, |
| keepdim: bool, |
| ) -> torch.Tensor: |
| var = inputs.mul(inputs).mean(normalization_axis, keepdim) |
| pre_shift_scale_norm_output = inputs / torch.sqrt(var + 1e-12) |
| norm_output = weight * pre_shift_scale_norm_output |
| return norm_output |
| |
| def nvfuser_fusion( |
| fd: FusionDefinition, |
| normalization_axis: int, |
| norm_size: int, |
| input_shape: List[int], |
| eps: float, |
| keepDim: bool |
| ) -> None : |
| inputs = fd.define_tensor(symbolic_sizes=[-1, -1, -1], contiguous=[True, True, True], dtype=DataType.Float) |
| weights = fd.define_tensor(symbolic_sizes=[-1], contiguous=[True], dtype=DataType.Float) |
| inputs_sq = fd.ops.mul(inputs, inputs) |
| sum0 = fd.ops.sum(inputs_sq, axes=[normalization_axis], keepdim=keepDim) |
| norm_const = fd.define_constant(norm_size) |
| var = fd.ops.div(sum0, norm_const) |
| eps_const = fd.define_constant(eps) |
| var_eps = fd.ops.add(var, eps_const) |
| invstd = fd.ops.rsqrt(var_eps) |
| pre_scale = fd.ops.mul(inputs, invstd) |
| weights_bcast = fd.ops.broadcast_in_dim(weights, output_shape=input_shape, broadcast_dims=[2]) |
| out = fd.ops.mul(pre_scale, weights_bcast) |
| fd.add_output(out) |
| fd.add_output(invstd) |
| |
| input_size = [64, 128, 1024] |
| dtype = torch.float32 |
| device = 'cuda' |
| inputs = torch.randn(*input_size, device=device, requires_grad=True) |
| weights = torch.nn.Parameter(torch.randn(input_size[2], dtype=dtype, device=device)) |
| fc = FusionCache.get() |
| before_fusions = fc.num_fusions() |
| |
| for _ in range(5) : |
| nvf_fusion = Fusion() |
| with FusionDefinition(nvf_fusion) as fd: |
| nvfuser_fusion(fd, 2, inputs.size()[2], inputs.size(), 1e-12, True) |
| nvf_out = nvf_fusion.execute([inputs, weights]) |
| |
| for _ in range(5) : |
| eager_out = primitive_definition(inputs, weights, 2, True) |
| |
| self.assertEqual(eager_out, nvf_out[0]) |
| self.assertEqual(fc.num_fusions() - before_fusions, 1) |
| |
| if __name__ == '__main__': |
| run_tests() |
