Star851
1340
代码
代码
Issues1541
Pull Requests2229
项目讨论
Wiki
分析
Star851
1340
已合并
feat support device parameters for aoti test #36115
huangyunlong创建于 5月19日
feat support device parameters for aoti test #36115
已合并
huangyunlong创建于 5月19日
从huangyunlong2022/pytorch2.10shim_td合入到Ascend/pytorchv2.10.0
共 2 个文件变更+114-40
Mtest/cpp_extensions/test/test_cpp_extensions_aot.py+22-22
@@ -921,14 +921,14 @@ class TestStableLibtorch(TestCase):
921 921 
922 def test_my_hann_window(self):922 def test_my_hann_window(self):
923 window_length = 5923 window_length = 5
924- dtype = None924+ dtype = torch.float32
925 layout = None925 layout = None
926- device = None926+ device = torch.device("npu")
927 pin_memory = None927 pin_memory = None
928 res = torch.ops.libtorch_agn_211.my_hann_window(928 res = torch.ops.libtorch_agn_211.my_hann_window(
929 window_length, dtype, layout, device, pin_memory929 window_length, dtype, layout, device, pin_memory
930 )930 )
931- expected = torch.hann_window(window_length)931+ expected = torch.hann_window(window_length, dtype=torch.float32, device="npu")
932 self.assertEqual(res, expected)932 self.assertEqual(res, expected)
933 933 
934 def test_my_histc(self):934 def test_my_histc(self):
@@ -1248,20 +1248,20 @@ class TestStableLibtorch(TestCase):
1248 size = [2, 3]1248 size = [2, 3]
1249 seed = 421249 seed = 42
1250 torch.manual_seed(seed)1250 torch.manual_seed(seed)
1251- res = torch.ops.libtorch_agn_211.my_rand(size, None, None, None, None)1251+ res = torch.ops.libtorch_agn_211.my_rand(size, None, None, torch.device("npu"), None)
1252 torch.manual_seed(seed)1252 torch.manual_seed(seed)
1253- expected = torch.rand(size)1253+ expected = torch.rand(size, device="npu")
1254- self.assertEqual(res.cpu(), expected)1254+ self.assertEqual(res.cpu(), expected.cpu())
1255 1255 
1256 def test_my_rand_generator(self):1256 def test_my_rand_generator(self):
1257 size = [2, 3]1257 size = [2, 3]
1258 seed = 421258 seed = 42
1259 torch.manual_seed(seed)1259 torch.manual_seed(seed)
1260 res = torch.ops.libtorch_agn_211.my_rand_generator(1260 res = torch.ops.libtorch_agn_211.my_rand_generator(
1261- size, None, None, None, None, None1261+ size, None, None, None, torch.device("npu"), None
1262 )1262 )
1263 torch.manual_seed(seed)1263 torch.manual_seed(seed)
1264- expected = torch.rand(size)1264+ expected = torch.rand(size, device="npu")
1265 self.assertEqual(res, expected)1265 self.assertEqual(res, expected)
1266 1266 
1267 def test_my_randint(self):1267 def test_my_randint(self):
@@ -1269,10 +1269,10 @@ class TestStableLibtorch(TestCase):
1269 size = [2, 3]1269 size = [2, 3]
1270 seed = 421270 seed = 42
1271 torch.manual_seed(seed)1271 torch.manual_seed(seed)
1272- res = torch.ops.libtorch_agn_211.my_randint(high, size, None, None, None, None)1272+ res = torch.ops.libtorch_agn_211.my_randint(high, size, None, None, torch.device("npu"), None)
1273 torch.manual_seed(seed)1273 torch.manual_seed(seed)
1274- expected = torch.randint(high, size)1274+ expected = torch.randint(high, size, device="npu")
1275- self.assertEqual(res.cpu(), expected)1275+ self.assertEqual(res.cpu(), expected.cpu())
1276 1276 
1277 def test_my_randint_generator(self):1277 def test_my_randint_generator(self):
1278 high = 101278 high = 10
@@ -1280,10 +1280,10 @@ class TestStableLibtorch(TestCase):
1280 seed = 421280 seed = 42
1281 torch.manual_seed(seed)1281 torch.manual_seed(seed)
1282 res = torch.ops.libtorch_agn_211.my_randint_generator(1282 res = torch.ops.libtorch_agn_211.my_randint_generator(
1283- high, size, None, None, None, None, None1283+ high, size, None, None, None, torch.device("npu"), None
1284 )1284 )
1285 torch.manual_seed(seed)1285 torch.manual_seed(seed)
1286- expected = torch.randint(high, size)1286+ expected = torch.randint(high, size, device="npu")
1287 self.assertEqual(res, expected)1287 self.assertEqual(res, expected)
1288 1288 
1289 def test_my_randint_low(self):1289 def test_my_randint_low(self):
@@ -1293,11 +1293,11 @@ class TestStableLibtorch(TestCase):
1293 seed = 421293 seed = 42
1294 torch.manual_seed(seed)1294 torch.manual_seed(seed)
1295 res = torch.ops.libtorch_agn_211.my_randint_low(1295 res = torch.ops.libtorch_agn_211.my_randint_low(
1296- low, high, size, None, None, None, None1296+ low, high, size, None, None, torch.device("npu"), None
1297 )1297 )
1298 torch.manual_seed(seed)1298 torch.manual_seed(seed)
1299- expected = torch.randint(low, high, size)1299+ expected = torch.randint(low, high, size, device="npu")
1300- self.assertEqual(res.cpu(), expected)1300+ self.assertEqual(res.cpu(), expected.cpu())
1301 1301 
1302 def test_my_randint_low_out(self):1302 def test_my_randint_low_out(self):
1303 out = torch.empty(2, 3, dtype=torch.long).npu()1303 out = torch.empty(2, 3, dtype=torch.long).npu()
@@ -1315,27 +1315,27 @@ class TestStableLibtorch(TestCase):
1315 size = [2, 3]1315 size = [2, 3]
1316 seed = 421316 seed = 42
1317 torch.manual_seed(seed)1317 torch.manual_seed(seed)
1318- res = torch.ops.libtorch_agn_211.my_randn(size, None, None, None, None)1318+ res = torch.ops.libtorch_agn_211.my_randn(size, None, None, torch.device("npu"), None)
1319 torch.manual_seed(seed)1319 torch.manual_seed(seed)
1320- expected = torch.randn(size)1320+ expected = torch.randn(size, device="npu")
1321- self.assertEqual(res.cpu(), expected)1321+ self.assertEqual(res.cpu(), expected.cpu())
1322 1322 
1323 def test_my_randn_generator(self):1323 def test_my_randn_generator(self):
1324 size = [2, 3]1324 size = [2, 3]
1325 seed = 421325 seed = 42
1326 torch.manual_seed(seed)1326 torch.manual_seed(seed)
1327 res = torch.ops.libtorch_agn_211.my_randn_generator(1327 res = torch.ops.libtorch_agn_211.my_randn_generator(
1328- size, None, None, None, None, None1328+ size, None, None, None, torch.device("npu"), None
1329 )1329 )
1330 torch.manual_seed(seed)1330 torch.manual_seed(seed)
1331- expected = torch.randn(size)1331+ expected = torch.randn(size, device="npu")
1332 self.assertEqual(res, expected)1332 self.assertEqual(res, expected)
1333 1333 
1334 def test_my_randperm(self):1334 def test_my_randperm(self):
1335 n = 51335 n = 5
1336 seed = 421336 seed = 42
1337 torch.npu.manual_seed(seed)1337 torch.npu.manual_seed(seed)
1338- res = torch.ops.libtorch_agn_211.my_randperm(n, None, None, None, None)1338+ res = torch.ops.libtorch_agn_211.my_randperm(n, None, None, torch.device("npu"), None)
1339 torch.npu.manual_seed(seed)1339 torch.npu.manual_seed(seed)
1340 expected = torch.randperm(n, device="npu")1340 expected = torch.randperm(n, device="npu")
1341 self.assertEqual(res.cpu(), expected.cpu())1341 self.assertEqual(res.cpu(), expected.cpu())
Mtest/cpp_extensions/test_stable_libtorch.cpp+92-18
@@ -1,8 +1,10 @@
1#include <torch/csrc/stable/library.h>1#include <torch/csrc/stable/library.h>
2#include <torch/csrc/stable/tensor.h>2#include <torch/csrc/stable/tensor.h>
3+#include <torch/csrc/stable/device.h>
3#include "torch_npu/csrc/inductor/aoti_torch/generated/c_shim_npu.h"4#include "torch_npu/csrc/inductor/aoti_torch/generated/c_shim_npu.h"
4 5 
5using torch::stable::Tensor;6using torch::stable::Tensor;
7+using torch::stable::Device;
6 8 
7Tensor my__adaptive_avg_pool2d(Tensor self, std::vector<int64_t> output_size)9Tensor my__adaptive_avg_pool2d(Tensor self, std::vector<int64_t> output_size)
8{10{
@@ -390,10 +392,18 @@ std::tuple<Tensor, Tensor> my_grid_sampler_2d_backward(Tensor grad_output, Tenso
390 return std::make_tuple(Tensor(ret0), Tensor(ret1));392 return std::make_tuple(Tensor(ret0), Tensor(ret1));
391}393}
392 394 
393-Tensor my_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)395+Tensor my_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
394{396{
395 AtenTensorHandle ret0;397 AtenTensorHandle ret0;
396- aoti_torch_npu_hann_window(window_length, dtype, layout, device, 0, pin_memory, &ret0);398+ int32_t device_val = 0;
399+ int32_t* device_ptr = nullptr;
400+ int32_t device_index = 0;
401+ if (device.has_value()) {
402+ device_val = static_cast<int32_t>(device.value().type());
403+ device_ptr = &device_val;
404+ device_index = device.value().index();
405+ }
406+ aoti_torch_npu_hann_window(window_length, dtype, layout, device_ptr, device_index, pin_memory, &ret0);
397 return Tensor(ret0);407 return Tensor(ret0);
398}408}
399 409 
@@ -611,40 +621,80 @@ Tensor my_pow_Tensor_Tensor(Tensor self, Tensor exponent)
611 return Tensor(ret0);621 return Tensor(ret0);
612}622}
613 623 
614-Tensor my_rand(std::vector<int64_t> size, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)624+Tensor my_rand(std::vector<int64_t> size, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
615{625{
616 AtenTensorHandle ret0;626 AtenTensorHandle ret0;
617- aoti_torch_npu_rand(size.data(), size.size(), dtype, layout, device, 0, pin_memory, &ret0);627+ int32_t device_val = 0;
628+ int32_t* device_ptr = nullptr;
629+ int32_t device_index = 0;
630+ if (device.has_value()) {
631+ device_val = static_cast<int32_t>(device.value().type());
632+ device_ptr = &device_val;
633+ device_index = device.value().index();
634+ }
635+ aoti_torch_npu_rand(size.data(), size.size(), dtype, layout, device_ptr, device_index, pin_memory, &ret0);
618 return Tensor(ret0);636 return Tensor(ret0);
619}637}
620 638 
621-Tensor my_rand_generator(std::vector<int64_t> size, std::optional<AtenGeneratorHandle> generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)639+Tensor my_rand_generator(std::vector<int64_t> size, std::optional<AtenGeneratorHandle> generator, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
622{640{
623 AtenTensorHandle ret0;641 AtenTensorHandle ret0;
624 AtenGeneratorHandle generator_handle = generator.value_or(nullptr);642 AtenGeneratorHandle generator_handle = generator.value_or(nullptr);
625- aoti_torch_npu_rand_generator(size.data(), size.size(), generator_handle ? &generator_handle : nullptr, dtype, layout, device, 0, pin_memory, &ret0);643+ int32_t device_val = 0;
644+ int32_t* device_ptr = nullptr;
645+ int32_t device_index = 0;
646+ if (device.has_value()) {
647+ device_val = static_cast<int32_t>(device.value().type());
648+ device_ptr = &device_val;
649+ device_index = device.value().index();
650+ }
651+ aoti_torch_npu_rand_generator(size.data(), size.size(), generator_handle ? &generator_handle : nullptr, dtype, layout, device_ptr, device_index, pin_memory, &ret0);
626 return Tensor(ret0);652 return Tensor(ret0);
627}653}
628 654 
629-Tensor my_randint(int64_t high, std::vector<int64_t> size, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)655+Tensor my_randint(int64_t high, std::vector<int64_t> size, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
630{656{
631 AtenTensorHandle ret0;657 AtenTensorHandle ret0;
632- aoti_torch_npu_randint(high, size.data(), size.size(), dtype, layout, device, 0, pin_memory, &ret0);658+ int32_t device_val = 0;
659+ int32_t* device_ptr = nullptr;
660+ int32_t device_index = 0;
661+ if (device.has_value()) {
662+ device_val = static_cast<int32_t>(device.value().type());
663+ device_ptr = &device_val;
664+ device_index = device.value().index();
665+ }
666+ aoti_torch_npu_randint(high, size.data(), size.size(), dtype, layout, device_ptr, device_index, pin_memory, &ret0);
633 return Tensor(ret0);667 return Tensor(ret0);
634}668}
635 669 
636-Tensor my_randint_generator(int64_t high, std::vector<int64_t> size, std::optional<AtenGeneratorHandle> generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)670+Tensor my_randint_generator(int64_t high, std::vector<int64_t> size, std::optional<AtenGeneratorHandle> generator, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
637{671{
638 AtenTensorHandle ret0;672 AtenTensorHandle ret0;
639 AtenGeneratorHandle generator_handle = generator.value_or(nullptr);673 AtenGeneratorHandle generator_handle = generator.value_or(nullptr);
640- aoti_torch_npu_randint_generator(high, size.data(), size.size(), generator_handle ? &generator_handle : nullptr, dtype, layout, device, 0, pin_memory, &ret0);674+ int32_t device_val = 0;
675+ int32_t* device_ptr = nullptr;
676+ int32_t device_index = 0;
677+ if (device.has_value()) {
678+ device_val = static_cast<int32_t>(device.value().type());
679+ device_ptr = &device_val;
680+ device_index = device.value().index();
681+ }
682+ aoti_torch_npu_randint_generator(high, size.data(), size.size(), generator_handle ? &generator_handle : nullptr, dtype, layout, device_ptr, device_index, pin_memory, &ret0);
641 return Tensor(ret0);683 return Tensor(ret0);
642}684}
643 685 
644-Tensor my_randint_low(int64_t low, int64_t high, std::vector<int64_t> size, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)686+Tensor my_randint_low(int64_t low, int64_t high, std::vector<int64_t> size, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
645{687{
646 AtenTensorHandle ret0;688 AtenTensorHandle ret0;
647- aoti_torch_npu_randint_low(low, high, size.data(), size.size(), dtype, layout, device, 0, pin_memory, &ret0);689+ int32_t device_val = 0;
690+ int32_t* device_ptr = nullptr;
691+ int32_t device_index = 0;
692+ if (device.has_value()) {
693+ device_val = static_cast<int32_t>(device.value().type());
694+ device_ptr = &device_val;
695+ device_index = device.value().index();
696+ }
697+ aoti_torch_npu_randint_low(low, high, size.data(), size.size(), dtype, layout, device_ptr, device_index, pin_memory, &ret0);
648 return Tensor(ret0);698 return Tensor(ret0);
649}699}
650 700 
@@ -653,25 +703,49 @@ void my_randint_low_out(Tensor out, int64_t low, int64_t high, std::vector<int64
653 aoti_torch_npu_randint_low_out(out.get(), low, high, size.data(), size.size());703 aoti_torch_npu_randint_low_out(out.get(), low, high, size.data(), size.size());
654}704}
655 705 
656-Tensor my_randn(std::vector<int64_t> size, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)706+Tensor my_randn(std::vector<int64_t> size, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
657{707{
658 AtenTensorHandle ret0;708 AtenTensorHandle ret0;
659- aoti_torch_npu_randn(size.data(), size.size(), dtype, layout, device, 0, pin_memory, &ret0);709+ int32_t device_val = 0;
710+ int32_t* device_ptr = nullptr;
711+ int32_t device_index = 0;
712+ if (device.has_value()) {
713+ device_val = static_cast<int32_t>(device.value().type());
714+ device_ptr = &device_val;
715+ device_index = device.value().index();
716+ }
717+ aoti_torch_npu_randn(size.data(), size.size(), dtype, layout, device_ptr, device_index, pin_memory, &ret0);
660 return Tensor(ret0);718 return Tensor(ret0);
661}719}
662 720 
663-Tensor my_randn_generator(std::vector<int64_t> size, std::optional<AtenGeneratorHandle> generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)721+Tensor my_randn_generator(std::vector<int64_t> size, std::optional<AtenGeneratorHandle> generator, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
664{722{
665 AtenTensorHandle ret0;723 AtenTensorHandle ret0;
666 AtenGeneratorHandle generator_handle = generator.value_or(nullptr);724 AtenGeneratorHandle generator_handle = generator.value_or(nullptr);
667- aoti_torch_npu_randn_generator(size.data(), size.size(), generator_handle ? &generator_handle : nullptr, dtype, layout, device, 0, pin_memory, &ret0);725+ int32_t device_val = 0;
726+ int32_t* device_ptr = nullptr;
727+ int32_t device_index = 0;
728+ if (device.has_value()) {
729+ device_val = static_cast<int32_t>(device.value().type());
730+ device_ptr = &device_val;
731+ device_index = device.value().index();
732+ }
733+ aoti_torch_npu_randn_generator(size.data(), size.size(), generator_handle ? &generator_handle : nullptr, dtype, layout, device_ptr, device_index, pin_memory, &ret0);
668 return Tensor(ret0);734 return Tensor(ret0);
669}735}
670 736 
671-Tensor my_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t* pin_memory)737+Tensor my_randperm(int64_t n, int32_t* dtype, int32_t* layout, std::optional<Device> device, int32_t* pin_memory)
672{738{
673 AtenTensorHandle ret0;739 AtenTensorHandle ret0;
674- aoti_torch_npu_randperm(n, dtype, layout, device, 0, pin_memory, &ret0);740+ int32_t device_val = 0;
741+ int32_t* device_ptr = nullptr;
742+ int32_t device_index = 0;
743+ if (device.has_value()) {
744+ device_val = static_cast<int32_t>(device.value().type());
745+ device_ptr = &device_val;
746+ device_index = device.value().index();
747+ }
748+ aoti_torch_npu_randperm(n, dtype, layout, device_ptr, device_index, pin_memory, &ret0);
675 return Tensor(ret0);749 return Tensor(ret0);
676}750}
677 751