已合并
fix: fix npugrphs backend bug and add indexput test case #31726
luochao60创建于 3月12日
fix: fix npugrphs backend bug and add indexput test case #31726
已合并
共 2 个文件变更+156-4
| @@ -0,0 +1,151 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch_npu | ||
| 3 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 4 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +class TestAclGraphSpecialOp(TestCase): | ||
| 8 | + | ||
| 9 | + def _fn_masked_assign_fwd(inp: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: | ||
| 10 | + """ | ||
| 11 | + Core pattern: tensor[bool_mask] = -1 | ||
| 12 | + Compiles to aten.index_put_ with boolean indices → cudagraph-unsafe. | ||
| 13 | + A relu op follows to verify that the fallback path runs normally. | ||
| 14 | + """ | ||
| 15 | + x = inp.clone() | ||
| 16 | + x[mask] = -1.0 # aten.index_put_ / bool indices → skip aclgraph | ||
| 17 | + out = torch.relu(x) # normal op: must execute correctly after fallback | ||
| 18 | + return out | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + def _fn_masked_assign_fwd_bwd(inp: torch.Tensor) -> torch.Tensor: | ||
| 22 | + """ | ||
| 23 | + Dynamic bool mask derived from input values, followed by a mul op. | ||
| 24 | + Returns a scalar so that .backward() can be called directly. | ||
| 25 | + The backward must produce correct gradients through the fallback path. | ||
| 26 | + """ | ||
| 27 | + mask = inp > 0 # bool mask whose shape depends on runtime values | ||
| 28 | + x = inp.clone() | ||
| 29 | + x[mask] = -1.0 # aten.index_put_ with bool indices | ||
| 30 | + out = x * 2.0 # normal op following the special op | ||
| 31 | + return out.sum() | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + def test_masked_assign_forward(self): | ||
| 35 | + """ | ||
| 36 | + Verify that tensor[bool_mask] = -1 followed by relu produces the | ||
| 37 | + same result whether executed eagerly or through torch.compile. | ||
| 38 | + | ||
| 39 | + The compiled graph contains aten.index_put_ with bool indices, which | ||
| 40 | + triggers the cudagraph-unsafe detection in _graph_tree.py. The graph | ||
| 41 | + is therefore skipped from ACL graph capture and runs in eager mode. | ||
| 42 | + """ | ||
| 43 | + shape = (8, 16) | ||
| 44 | + npu = torch.device("npu") | ||
| 45 | + inp = torch.randn(shape, dtype=torch.float32, device=npu) | ||
| 46 | + mask = torch.randint(0, 2, shape, device=npu).bool() | ||
| 47 | + | ||
| 48 | + std_result = self._fn_masked_assign_fwd(inp, mask) | ||
| 49 | + | ||
| 50 | + compiled_fn = torch.compile( | ||
| 51 | + self._fn_masked_assign_fwd, backend="inductor", | ||
| 52 | + options={"triton.cudagraphs": True}, | ||
| 53 | + ) | ||
| 54 | + compiled_result = compiled_fn(inp, mask) | ||
| 55 | + | ||
| 56 | + self.assertEqual(std_result, compiled_result, prec=1e-3) | ||
| 57 | + | ||
| 58 | + | ||
| 59 | + def test_masked_assign_forward_backward(self): | ||
| 60 | + """ | ||
| 61 | + Verify forward output and input gradients for a function that: | ||
| 62 | + 1. Builds a dynamic bool mask (inp > 0) | ||
| 63 | + 2. Assigns -1 to masked positions (index_put_ with bool indices) | ||
| 64 | + 3. Multiplies the result by 2.0 (normal op) | ||
| 65 | + 4. Reduces to a scalar (sum) | ||
| 66 | + | ||
| 67 | + The compiled version must: | ||
| 68 | + - Detect the boolean index_put and skip ACL graph capture. | ||
| 69 | + - Fall back to eager, preserving correct autograd behaviour. | ||
| 70 | + - Produce the same scalar output and input gradients as eager. | ||
| 71 | + """ | ||
| 72 | + shape = (8, 16) | ||
| 73 | + npu = torch.device("npu") | ||
| 74 | + | ||
| 75 | + # Eager reference tensors | ||
| 76 | + x_ref = torch.randn(shape, dtype=torch.float32, device=npu, | ||
| 77 | + requires_grad=True) | ||
| 78 | + # Compiled-path tensors (same initial values) | ||
| 79 | + x = x_ref.detach().clone().requires_grad_(True) | ||
| 80 | + | ||
| 81 | + # ---- Eager forward + backward ---- | ||
| 82 | + std_out = self._fn_masked_assign_fwd_bwd(x_ref) | ||
| 83 | + std_out.backward() | ||
| 84 | + | ||
| 85 | + # ---- Compiled forward + backward ---- | ||
| 86 | + compiled_fn = torch.compile( | ||
| 87 | + self._fn_masked_assign_fwd_bwd, backend="inductor", | ||
| 88 | + options={"triton.cudagraphs": True}, | ||
| 89 | + ) | ||
| 90 | + compiled_out = compiled_fn(x) | ||
| 91 | + compiled_out.backward() | ||
| 92 | + | ||
| 93 | + # Forward output must match | ||
| 94 | + self.assertEqual(std_out, compiled_out, prec=1e-3) | ||
| 95 | + # Input gradients must match | ||
| 96 | + self.assertEqual(x.grad, x_ref.grad, prec=1e-3) | ||
| 97 | + | ||
| 98 | + | ||
| 99 | + | ||
| 100 | + def test_npugraphs_masked_assign_forward(self): | ||
| 101 | + """ | ||
| 102 | + Same as test_masked_assign_forward but using the npugraphs backend. | ||
| 103 | + Verifies that the npugraphs backend also correctly detects | ||
| 104 | + aten.index_put_ with bool indices and falls back to eager. | ||
| 105 | + """ | ||
| 106 | + shape = (8, 16) | ||
| 107 | + npu = torch.device("npu") | ||
| 108 | + inp = torch.randn(shape, dtype=torch.float32, device=npu) | ||
| 109 | + mask = torch.randint(0, 2, shape, device=npu).bool() | ||
| 110 | + | ||
| 111 | + std_result = self._fn_masked_assign_fwd(inp, mask) | ||
| 112 | + | ||
| 113 | + compiled_fn = torch.compile( | ||
| 114 | + self._fn_masked_assign_fwd, backend="npugraphs", | ||
| 115 | + ) | ||
| 116 | + compiled_result = compiled_fn(inp, mask) | ||
| 117 | + | ||
| 118 | + self.assertEqual(std_result, compiled_result, prec=1e-3) | ||
| 119 | + | ||
| 120 | + | ||
| 121 | + def test_npugraphs_masked_assign_forward_backward(self): | ||
| 122 | + """ | ||
| 123 | + Same as test_masked_assign_forward_backward but using the npugraphs | ||
| 124 | + backend. Verifies forward output and input gradients are correct | ||
| 125 | + when the npugraphs backend falls back to eager due to bool index_put. | ||
| 126 | + """ | ||
| 127 | + shape = (8, 16) | ||
| 128 | + npu = torch.device("npu") | ||
| 129 | + | ||
| 130 | + x_ref = torch.randn(shape, dtype=torch.float32, device=npu, | ||
| 131 | + requires_grad=True) | ||
| 132 | + x = x_ref.detach().clone().requires_grad_(True) | ||
| 133 | + | ||
| 134 | + # ---- Eager forward + backward ---- | ||
| 135 | + std_out = self._fn_masked_assign_fwd_bwd(x_ref) | ||
| 136 | + std_out.backward() | ||
| 137 | + | ||
| 138 | + # ---- Compiled forward + backward ---- | ||
| 139 | + compiled_fn = torch.compile( | ||
| 140 | + self._fn_masked_assign_fwd_bwd, backend="npugraphs", | ||
| 141 | + ) | ||
| 142 | + compiled_out = compiled_fn(x) | ||
| 143 | + compiled_out.backward() | ||
| 144 | + | ||
| 145 | + self.assertEqual(std_out, compiled_out, prec=1e-3) | ||
| 146 | + self.assertEqual(x.grad, x_ref.grad, prec=1e-3) | ||
| 147 | + | ||
| 148 | + | ||
| 149 | + | ||
| 150 | +if __name__ == "__main__": | ||
| 151 | + run_tests() | ||
| @@ -186,12 +186,13 @@ def patch_get_first_incompatible_cudagraph_node(): | |||
| 186 | inductor_utils.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node | 186 | inductor_utils.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node |
| 187 | 187 | ||
| 188 | from torch._inductor import compile_fx | 188 | from torch._inductor import compile_fx |
| 189 | - if hasattr(compile_fx, 'get_first_incompatible_cudagraph_node'): | 189 | + compile_fx.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node |
| 190 | - compile_fx.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node | ||
| 191 | 190 | ||
| 192 | from torch._dynamo.backends import cudagraphs | 191 | from torch._dynamo.backends import cudagraphs |
| 193 | - if hasattr(cudagraphs, 'get_first_incompatible_cudagraph_node'): | 192 | + cudagraphs.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node |
| 194 | - cudagraphs.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node | 193 | + |
| 194 | + from torch_npu.utils import _graph_tree | ||
| 195 | + _graph_tree.get_first_incompatible_cudagraph_node = get_first_incompatible_cudagraph_node | ||
| 195 | 196 | ||
| 196 | 197 | ||
| 197 | def disable_foreach(): | 198 | def disable_foreach(): |