已合并
meta_check #29835
cuiduo创建于 1月21日
meta_check #29835
已合并
共 6 个文件变更+110-9
| @@ -0,0 +1,30 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + parametrize, | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | +) | ||
| 7 | +from testutils import TestUtils | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestResizeAs(TestUtils): | ||
| 12 | + def op_calc(self, x, target, adder): | ||
| 13 | + return torch.ops.aten.resize_as_.default(x, target) + adder | ||
| 14 | + | ||
| 15 | + def test_default(self): | ||
| 16 | + shape = (2, 8) | ||
| 17 | + target_shape = (8, 2) | ||
| 18 | + x = torch.randn(shape, dtype=torch.float16, device="npu") | ||
| 19 | + target = torch.randn(target_shape, dtype=torch.float16, device="npu") | ||
| 20 | + adder = torch.randn(target_shape, dtype=torch.float16, device="npu") | ||
| 21 | + std_out = self.op_calc(x, target, adder) | ||
| 22 | + compile_out = torch.compile(self.op_calc, backend='inductor')(x, target, adder) | ||
| 23 | + self.assertEqual(std_out, compile_out, atol=1e-5, rtol=1e-5) | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +instantiate_parametrized_tests(TestResizeAs) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +if __name__ == "__main__": | ||
| 30 | + run_tests() | ||
| @@ -0,0 +1,40 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + parametrize, | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | +) | ||
| 7 | +from testutils import TestUtils | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestSqueeze(TestUtils): | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + def test_squeeze_cases(self, shape, dtype): | ||
| 16 | + def op_calc(x): | ||
| 17 | + return torch.ops.aten.squeeze_.default(x) | ||
| 18 | + | ||
| 19 | + tensor = self._generate_tensor(shape, dtype) | ||
| 20 | + std_result = op_calc(tensor) | ||
| 21 | + compile_result = torch.compile(op_calc, backend='inductor')(tensor) | ||
| 22 | + self.assertEqual(std_result, compile_result, atol=1e-5, rtol=1e-5) | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + def test_squeeze_dims_cases(self, shape, dtype): | ||
| 27 | + def op_calc(x): | ||
| 28 | + return torch.ops.aten.squeeze_.dims(x, dim=[0]) | ||
| 29 | + | ||
| 30 | + tensor = self._generate_tensor(shape, dtype) | ||
| 31 | + std_result = op_calc(tensor) | ||
| 32 | + compile_result = torch.compile(op_calc, backend='inductor')(tensor) | ||
| 33 | + self.assertEqual(std_result, compile_result, atol=1e-5, rtol=1e-5) | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +instantiate_parametrized_tests(TestSqueeze) | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +if __name__ == "__main__": | ||
| 40 | + run_tests() | ||
| @@ -44,9 +44,15 @@ const at::Tensor& NPUNativeFunctions::resize_( | |||
| 44 | } | 44 | } |
| 45 | // because of resize _impl_npu_ only support at base format, so | 45 | // because of resize _impl_npu_ only support at base format, so |
| 46 | // no need to reflush NpuStorageDesc here. | 46 | // no need to reflush NpuStorageDesc here. |
| 47 | - at::Tensor temp_self = self; | 47 | + auto ks = self.key_set(); |
| 48 | - if (!FormatHelper::IsBaseFormatType(self)) { | 48 | + bool is_fake_or_meta = ks.has_all(c10::DispatchKeySet(c10::BackendComponent::MetaBit)) || |
| 49 | - NPUNativeFunctions::npu_format_cast_(temp_self, FormatHelper::GetBaseFormat(self)); | 49 | + ks.has_all(c10::DispatchKeySet(c10::DispatchKey::Python)) || |
| 50 | + self.is_meta(); | ||
| 51 | + if (!is_fake_or_meta) { | ||
| 52 | + at::Tensor temp_self = self; | ||
| 53 | + if (!FormatHelper::IsBaseFormatType(self)) { | ||
| 54 | + NPUNativeFunctions::npu_format_cast_(temp_self, FormatHelper::GetBaseFormat(self)); | ||
| 55 | + } | ||
| 50 | } | 56 | } |
| 51 | auto* self_ = self.unsafeGetTensorImpl(); | 57 | auto* self_ = self.unsafeGetTensorImpl(); |
| 52 | resize_impl_npu_(self_, size, c10::nullopt); | 58 | resize_impl_npu_(self_, size, c10::nullopt); |
| @@ -144,7 +144,7 @@ inline at::TensorImpl* resize_impl_npu_( | |||
| 144 | storage_size = self->numel(); | 144 | storage_size = self->numel(); |
| 145 | } | 145 | } |
| 146 | maybe_resize_storage_npu(self, storage_size, size); | 146 | maybe_resize_storage_npu(self, storage_size, size); |
| 147 | - | 147 | + |
 | |||
| 148 | return self; | 148 | return self; |
| 149 | } | 149 | } |
| 150 | 150 | ||
| @@ -146,11 +146,17 @@ const at::Tensor& NPUNativeFunctions::as_strided__symint( | |||
| 146 | c10::SymIntArrayRef stride, | 146 | c10::SymIntArrayRef stride, |
| 147 | c10::optional<c10::SymInt> storage_offset_) | 147 | c10::optional<c10::SymInt> storage_offset_) |
| 148 | { | 148 | { |
| 149 | - if (InferFormat::IsDefiniteTensorWhenMetaDataChanges(self, c10::asIntArrayRefUnchecked(size)) && | 149 | + auto ks = self.key_set(); |
| 150 | - !FormatHelper::IsOpInputBaseFormat(self)) { | 150 | + bool is_fake_or_meta = ks.has_all(c10::DispatchKeySet(c10::BackendComponent::MetaBit)) || |
| 151 | - TORCH_CHECK(false, "Current tensor is running as_strided__symint while internal format is not allowed." | 151 | + ks.has_all(c10::DispatchKeySet(c10::DispatchKey::Python)) || |
| 152 | - " You can try torch.npu.config.allow_internal_format = False to avoid the problem.", | 152 | + self.is_meta(); |
| 153 | - PTA_ERROR(ErrCode::NOT_SUPPORT)) | 153 | + if (!is_fake_or_meta) { |
| 154 | + if (InferFormat::IsDefiniteTensorWhenMetaDataChanges(self, c10::asIntArrayRefUnchecked(size)) && | ||
| 155 | + !FormatHelper::IsOpInputBaseFormat(self)) { | ||
| 156 | + TORCH_CHECK(false, "Current tensor is running as_strided__symint while internal format is not allowed." | ||
| 157 | + " You can try torch.npu.config.allow_internal_format = False to avoid the problem.", | ||
| 158 | + PTA_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 159 | + } | ||
| 154 | } | 160 | } |
| 155 | auto storage_offset = storage_offset_.value_or(self.sym_storage_offset()); | 161 | auto storage_offset = storage_offset_.value_or(self.sym_storage_offset()); |
| 156 | at::native::setStrided(self, size, stride, std::move(storage_offset)); | 162 | at::native::setStrided(self, size, stride, std::move(storage_offset)); |
| @@ -236,6 +236,24 @@ def patch_dynamo_optimize(): | |||
| 236 | torch._dynamo.optimize = npu_optimize | 236 | torch._dynamo.optimize = npu_optimize |
| 237 | 237 | ||
| 238 | 238 | ||
| 239 | +def patch_base_schedulernode(): | ||
| 240 | + from torch._inductor.scheduler import BaseSchedulerNode | ||
| 241 | + from torch._inductor.scheduler import ExternKernelSchedulerNode | ||
| 242 | + | ||
| 243 | + original_get_read_write_buffer_accesses = BaseSchedulerNode.get_read_write_buffer_accesses | ||
| 244 | + | ||
| 245 | + def new_get_read_write_buffer_accesses( | ||
| 246 | + self_instance, include_reads: bool, include_writes: bool | ||
| 247 | + ) -> dict[str, int]: | ||
| 248 | + if isinstance(self_instance, ExternKernelSchedulerNode): | ||
| 249 | + return {} | ||
| 250 | + return original_get_read_write_buffer_accesses( | ||
| 251 | + self_instance, include_reads, include_writes | ||
| 252 | + ) | ||
| 253 | + | ||
| 254 | + BaseSchedulerNode.get_read_write_buffer_accesses = new_get_read_write_buffer_accesses | ||
| 255 | + | ||
| 256 | + | ||
| 239 | def add_dynamo_methods(): | 257 | def add_dynamo_methods(): |
| 240 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ | 258 | UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__ |
| 241 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ | 259 | UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__ |
| @@ -245,3 +263,4 @@ def add_dynamo_methods(): | |||
| 245 | TensorVariable.call_method = TensorVariable_call_method | 263 | TensorVariable.call_method = TensorVariable_call_method |
| 246 | patch_dynamo_optimize() | 264 | patch_dynamo_optimize() |
| 247 | patch_inductor_wrapper() | 265 | patch_inductor_wrapper() |
| 266 | + patch_base_schedulernode() | ||
删掉多余的空格