已合并
[fix] restore test_storage_size usecase #37127
maoyuanpeng1创建于 5月29日
[fix] restore test_storage_size usecase #37127
已合并
共 1 个文件变更+12-0
| @@ -44,6 +44,18 @@ class TestCppExtensionAOT(TestCase): | |||
| 44 | npu_z = npu_extension.npu_add(x.npu(), y.npu()) | 44 | npu_z = npu_extension.npu_add(x.npu(), y.npu()) |
| 45 | self.assertEqual(npu_z.cpu(), (x + y)) | 45 | self.assertEqual(npu_z.cpu(), (x + y)) |
| 46 | 46 | ||
| 47 | + def test_storage_sizes(self): | ||
| 48 | + torch.npu.config.allow_internal_format = True | ||
| 49 | + t = torch_npu.npu_format_cast(torch.ones(128, 512, dtype=torch.int8).npu(), 29) | ||
| 50 | + self.assertTrue(npu_extension.check_storage_sizes(t, (16, 8, 16, 32))) | ||
| 51 | + t = torch_npu.npu_format_cast(torch.ones(31, 127, 511, dtype=torch.int8).npu(), 29) | ||
| 52 | + self.assertTrue(npu_extension.check_storage_sizes(t, (31, 16, 8, 16, 32))) | ||
| 53 | + t = torch_npu.npu_format_cast(torch.ones(128, 512, dtype=torch.float16).npu(), 29) | ||
| 54 | + self.assertTrue(npu_extension.check_storage_sizes(t, (32, 8, 16, 16))) | ||
| 55 | + # float32 will cast to float16 before calculate | ||
| 56 | + t = torch_npu.npu_format_cast(torch.ones(128, 512, dtype=torch.float32).npu(), 29) | ||
| 57 | + self.assertTrue(npu_extension.check_storage_sizes(t, (32, 8, 16, 16))) | ||
| 58 | + | ||
| 47 | def test_dispatch_allreduce(self): | 59 | def test_dispatch_allreduce(self): |
| 48 | flags = os.O_WRONLY | os.O_RDONLY | os.O_CREAT | 60 | flags = os.O_WRONLY | os.O_RDONLY | os.O_CREAT |
| 49 | modes = stat.S_IWUSR | stat.S_IRUSR | 61 | modes = stat.S_IWUSR | stat.S_IRUSR |