已合并
fix: [910B/910_93] fall back to aclop for view tensors in format cast #45058
wuyouqi1创建于 15 天前
fix: [910B/910_93] fall back to aclop for view tensors in format cast #45058
已合并
共 2 个文件变更+74-11
| @@ -8,7 +8,9 @@ from torch_npu.testing.common_utils import create_common_tensor, SupportedDevice | |||
| 8 | torch.npu.config.allow_internal_format = True | 8 | torch.npu.config.allow_internal_format = True |
| 9 | 9 | ||
| 10 | # ACL format constants | 10 | # ACL format constants |
| 11 | -ACL_FORMAT_ND = 2 | 11 | +ACL_FORMAT_NCHW = 0 |
| 12 | +ACL_FORMAT_ND = 2 | ||
| 13 | +ACL_FORMAT_NC1HWC0 = 3 | ||
| 12 | ACL_FORMAT_FRACTAL_NZ = 29 | 14 | ACL_FORMAT_FRACTAL_NZ = 29 |
| 13 | 15 | ||
| 14 | 16 | ||
| @@ -251,6 +253,61 @@ class TestNpuFormatCastAclnn(TestCase): | |||
| 251 | self.assertEqual(out.shape, (N, 1)) | 253 | self.assertEqual(out.shape, (N, 1)) |
| 252 | self.assertEqual(out.storage()[1], t_view.storage()[1]) | 254 | self.assertEqual(out.storage()[1], t_view.storage()[1]) |
| 253 | 255 | ||
| 256 | + # ------------------------------------------------------------------ # | ||
| 257 | + # Group 8: View guard — sizes mismatch with storage desc falls back | ||
| 258 | + # to aclop (910B/910_93) | ||
| 259 | + # ------------------------------------------------------------------ # | ||
| 260 | + | ||
| 261 | + | ||
| 262 | + def test_view_reshape_internal_to_base_fallback(self): | ||
| 263 | + """NC1HWC0 -> reshape view -> NCHW: falls back to aclop (no OOM).""" | ||
| 264 | + t = torch_npu.npu_format_cast(torch.rand(4, 8, 28, 28).float().npu(), ACL_FORMAT_NC1HWC0) | ||
| 265 | + v = t.reshape(4, -1) | ||
| 266 | + out = torch_npu.npu_format_cast(v, ACL_FORMAT_NCHW) | ||
| 267 | + self.assertEqual(torch_npu.get_npu_format(out), ACL_FORMAT_NCHW) | ||
| 268 | + self.assertTrue(torch.allclose(out.cpu(), v.cpu(), rtol=1e-3, atol=1e-5)) | ||
| 269 | + | ||
| 270 | + | ||
| 271 | + def test_view_reshape_nd_to_nz_fallback(self): | ||
| 272 | + """ND -> reshape view -> NZ (fp16): falls back to aclop (no OOM).""" | ||
| 273 | + t = torch.rand(32, 64).half().npu() | ||
| 274 | + v = t.reshape(-1) | ||
| 275 | + out = torch_npu.npu_format_cast(v, ACL_FORMAT_FRACTAL_NZ) | ||
| 276 | + self.assertEqual(torch_npu.get_npu_format(out), ACL_FORMAT_FRACTAL_NZ) | ||
| 277 | + back = torch_npu.npu_format_cast(out, ACL_FORMAT_ND) | ||
| 278 | + self.assertTrue(torch.allclose(back.cpu(), v.cpu(), rtol=1e-2, atol=1e-3)) | ||
| 279 | + | ||
| 280 | + | ||
| 281 | + def test_view_transpose_nd_to_nz_fallback(self): | ||
| 282 | + """ND -> transpose view (non-contiguous) -> NZ: falls back to aclop.""" | ||
| 283 | + t = torch.rand(32, 64).half().npu() | ||
| 284 | + v = t.t() | ||
| 285 | + out = torch_npu.npu_format_cast(v, ACL_FORMAT_FRACTAL_NZ) | ||
| 286 | + self.assertEqual(torch_npu.get_npu_format(out), ACL_FORMAT_FRACTAL_NZ) | ||
| 287 | + back = torch_npu.npu_format_cast(out, ACL_FORMAT_ND) | ||
| 288 | + self.assertTrue(torch.allclose(back.cpu(), v.cpu(), rtol=1e-2, atol=1e-3)) | ||
| 289 | + | ||
| 290 | + | ||
| 291 | + def test_view_strided_same_shape_nd_to_nz(self): | ||
| 292 | + """Same-shape strided view (as_strided) -> NZ: handled on aclnn.""" | ||
| 293 | + t = torch.rand(32, 64).half().npu() | ||
| 294 | + v = torch.as_strided(t, (32, 64), (1, 32)) | ||
| 295 | + self.assertEqual(v.shape, t.shape) | ||
| 296 | + out = torch_npu.npu_format_cast(v, ACL_FORMAT_FRACTAL_NZ) | ||
| 297 | + self.assertEqual(torch_npu.get_npu_format(out), ACL_FORMAT_FRACTAL_NZ) | ||
| 298 | + back = torch_npu.npu_format_cast(out, ACL_FORMAT_ND) | ||
| 299 | + self.assertTrue(torch.allclose(back.cpu(), v.cpu(), rtol=1e-2, atol=1e-3)) | ||
| 300 | + | ||
| 301 | + | ||
| 302 | + def test_view_strided_same_shape_internal_to_base(self): | ||
| 303 | + """Same-shape strided view on internal format -> NCHW: aclop fallback.""" | ||
| 304 | + t = torch_npu.npu_format_cast(torch.rand(4, 8, 28, 28).float().npu(), ACL_FORMAT_NC1HWC0) | ||
| 305 | + v = torch.as_strided(t, (4, 8, 28, 28), (1, 4 * 28 * 28, 4 * 28, 4)) | ||
| 306 | + self.assertEqual(v.shape, t.shape) | ||
| 307 | + out = torch_npu.npu_format_cast(v, ACL_FORMAT_NCHW) | ||
| 308 | + self.assertEqual(torch_npu.get_npu_format(out), ACL_FORMAT_NCHW) | ||
| 309 | + self.assertTrue(torch.allclose(out.cpu(), v.cpu(), rtol=1e-3, atol=1e-5)) | ||
| 310 | + | ||
| 254 | 311 | ||
| 255 | class TestNpuFormatCastDtypeParam(TestCase): | 312 | class TestNpuFormatCastDtypeParam(TestCase): |
| 256 | """ | 313 | """ |
| @@ -70,6 +70,21 @@ static bool IsNzC0Variant(int64_t fmt) | |||
| 70 | fmt == ACL_FORMAT_FRACTAL_NZ_C0_8; | 70 | fmt == ACL_FORMAT_FRACTAL_NZ_C0_8; |
| 71 | } | 71 | } |
| 72 | 72 | ||
| 73 | +// Views (sizes mismatch with storage desc) fall back to aclop. | ||
| 74 | +static bool IsViewTensor(const at::Tensor& src) | ||
| 75 | +{ | ||
| 76 | + const auto& desc = torch_npu::NPUBridge::GetNpuStorageImpl(src)->npu_desc_; | ||
| 77 | + if (src.sizes().size() != static_cast<int64_t>(desc.base_sizes_.size())) { | ||
| 78 | + return true; | ||
| 79 | + } | ||
| 80 | + for (size_t i = 0; i < src.sizes().size(); ++i) { | ||
| 81 | + if (src.sizes()[i] != desc.base_sizes_[i]) { | ||
| 82 | + return true; | ||
| 83 | + } | ||
| 84 | + } | ||
| 85 | + return false; | ||
| 86 | +} | ||
| 87 | + | ||
| 73 | std::tuple<bool, int64_t, c10::SmallVector<int64_t, SIZE>> MaybeUseAclnnNpuFormatCast(const at::Tensor& src, | 88 | std::tuple<bool, int64_t, c10::SmallVector<int64_t, SIZE>> MaybeUseAclnnNpuFormatCast(const at::Tensor& src, |
| 74 | int64_t acl_format, c10::optional<int64_t> customize_dtype, c10::optional<int64_t> input_dtype) | 89 | int64_t acl_format, c10::optional<int64_t> customize_dtype, c10::optional<int64_t> input_dtype) |
| 75 | { | 90 | { |
| @@ -138,16 +153,7 @@ std::tuple<bool, int64_t, c10::SmallVector<int64_t, SIZE>> MaybeUseAclnnNpuForma | |||
| 138 | if (!FormatHelper::IsBaseFormatType(src) && !src.is_contiguous()) { | 153 | if (!FormatHelper::IsBaseFormatType(src) && !src.is_contiguous()) { |
| 139 | return std::make_tuple(false, dstFormat, outputShape); | 154 | return std::make_tuple(false, dstFormat, outputShape); |
| 140 | } | 155 | } |
| 141 | - // Guard: when casting a 2D contiguous tensor from ND to FRACTAL_NZ, if the | 156 | + if (IsViewTensor(src)) { |
| 142 | - // tensor's current shape differs from its storage shape (base_sizes_), it means | ||
| 143 | - // it is a view (e.g. transpose). Downstream ops may encounter precision | ||
| 144 | - // issues, so the aclnn path is not supported for this case. | ||
| 145 | - // Fall back to the aclop path to stay safe. | ||
| 146 | - // Example: storage [1, N] viewed as [N, 1] (Nx1 column vector). | ||
| 147 | - auto src_desc = torch_npu::NPUBridge::GetNpuStorageImpl(src)->npu_desc_; | ||
| 148 | - if (src_desc.npu_format_ == ACL_FORMAT_ND && acl_format == ACL_FORMAT_FRACTAL_NZ && | ||
| 149 | - src.is_contiguous() && src.sizes().size() == 2 && src_desc.base_sizes_.size() == 2 && src.sizes()[1] == 1 && | ||
| 150 | - (src_desc.base_sizes_[0] != src.sizes()[0] || src_desc.base_sizes_[1] != src.sizes()[1])) { | ||
| 151 | return std::make_tuple(false, dstFormat, outputShape); | 157 | return std::make_tuple(false, dstFormat, outputShape); |
| 152 | } | 158 | } |
| 153 | if (IsAclnnFormatCastSupported() && aclnnNpuFormatCastExist) { | 159 | if (IsAclnnFormatCastSupported() && aclnnNpuFormatCastExist) { |