fix: support to_sparse on non-default NPU devices-pytorch-AtomGit

已合并

fix: support to_sparse on non-default NPU devices #35947

hz893创建于 5月18日从hz893/pytorch： fix_to_sparse_v2.12.0 合入到 Ascend/pytorch： v2.12.0

fix: support to_sparse on non-default NPU devices #35947

已合并

hz893创建于 5月18日

从hz893/pytorch fix_to_sparse_v2.12.0合入到Ascend/pytorch v2.12.0

共 4 个文件变更+51-0

Mtest/test_sparse_coo.py+22-0

	@@ -610,6 +610,28 @@ class TestSparseCoo(TestCase):
610	self.assertEqual(res_npu.indices(), res_cpu.indices())	610	self.assertEqual(res_npu.indices(), res_cpu.indices())
611	self.assertEqual(res_npu.values(), res_cpu.values())	611	self.assertEqual(res_npu.values(), res_cpu.values())
612		612
		613	+ @unittest.skipIf(torch.npu.device_count() < 2, "only one NPU detected")
		614	+ @SupportedDevices(['Ascend910B'])
		615	+ def test_sparse_to_sparse_non_default_npu_device(self):
		616	+ original_device = torch.npu.current_device()
		617	+ try:
		618	+ torch.npu.set_device(0)
		619	+ tensor_npu = torch.tensor(
		620	+ [[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]], device="npu:1")
		621	+
		622	+ for sparse_dim in (None, 2):
		623	+ res_npu = (
		624	+ tensor_npu.to_sparse()
		625	+ if sparse_dim is None
		626	+ else tensor_npu.to_sparse(sparse_dim)
		627	+ )
		628	+ self.assertEqual(res_npu.device, torch.device("npu:1"))
		629	+ self.assertEqual(res_npu._indices().device, torch.device("npu:1"))
		630	+ self.assertEqual(res_npu._values().device, torch.device("npu:1"))
		631	+ self.assertRtolEqual(res_npu.to_dense(), tensor_npu)
		632	+ finally:
		633	+ torch.npu.set_device(original_device)
		634	+
613	@SupportedDevices(['Ascend910B'])	635	@SupportedDevices(['Ascend910B'])
614	def test_sparse_any(self):	636	def test_sparse_any(self):
615	sparse_coo_npu, sparse_coo_cpu = self._create_sparse_coo_coalesced_tensor(torch.float)	637	sparse_coo_npu, sparse_coo_cpu = self._create_sparse_coo_coalesced_tensor(torch.float)

Atorch_npu/csrc/aten/common/TensorConversions.cpp+26-0

@@ -0,0 +1,26 @@
	1	+#include <ATen/ATen.h>
	2	+#include <ATen/ops/_to_sparse_native.h>
	3	+
	4	+#include <optional>
	5	+
	6	+#include "torch_npu/csrc/aten/NPUNativeFunctions.h"
	7	+
	8	+namespace at_npu {
	9	+namespace native {
	10	+
	11	+at::Tensor NPUNativeFunctions::_to_sparse(const at::Tensor& self, int64_t sparse_dim)
	12	+{
	13	+ return at::native::dense_to_sparse(self, sparse_dim);
	14	+}
	15	+
	16	+at::Tensor NPUNativeFunctions::_to_sparse(
	17	+ const at::Tensor& self,
	18	+ ::std::optional<at::Layout> layout,
	19	+ at::OptionalIntArrayRef blocksize,
	20	+ ::std::optional<int64_t> dense_dim)
	21	+{
	22	+ return at::native::dense_to_sparse(self, layout, blocksize, dense_dim);
	23	+}
	24	+
	25	+} // namespace native
	26	+} // namespace at_npu

Mtorch_npu/csrc/aten/npu_native_functions.yaml+2-0

	@@ -46,6 +46,8 @@ supported:
46	- squeeze	46	- squeeze
47	- squeeze.dim	47	- squeeze.dim
48	- _to_copy	48	- _to_copy
		49	+ - _to_sparse
		50	+ - _to_sparse.sparse_dim
49	- tril_indices	51	- tril_indices
50	- triu_indices	52	- triu_indices
51	- unfold	53	- unfold

Mtorch_npu/csrc/inductor/aoti_torch/generated/c_shim_npu.h+1-0

	@@ -29,6 +29,7 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__pdist_forward(AtenTensorHandle
29	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);	29	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
30	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);	30	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
31	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__thnn_fused_lstm_cell(AtenTensorHandle input_gates, AtenTensorHandle hidden_gates, AtenTensorHandle cx, AtenTensorHandle* input_bias, AtenTensorHandle* hidden_bias, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);	31	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__thnn_fused_lstm_cell(AtenTensorHandle input_gates, AtenTensorHandle hidden_gates, AtenTensorHandle cx, AtenTensorHandle* input_bias, AtenTensorHandle* hidden_bias, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
		32	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
32	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__trilinear(AtenTensorHandle i1, AtenTensorHandle i2, AtenTensorHandle i3, const int64_t* expand1, int64_t expand1_len_, const int64_t* expand2, int64_t expand2_len_, const int64_t* expand3, int64_t expand3_len_, const int64_t* sumdim, int64_t sumdim_len_, int64_t unroll_dim, AtenTensorHandle* ret0);	33	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__trilinear(AtenTensorHandle i1, AtenTensorHandle i2, AtenTensorHandle i3, const int64_t* expand1, int64_t expand1_len_, const int64_t* expand2, int64_t expand2_len_, const int64_t* expand3, int64_t expand3_len_, const int64_t* sumdim, int64_t sumdim_len_, int64_t unroll_dim, AtenTensorHandle* ret0);
33	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_abs(AtenTensorHandle self, AtenTensorHandle* ret0);	34	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_abs(AtenTensorHandle self, AtenTensorHandle* ret0);
34	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_adaptive_max_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);	35	AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_adaptive_max_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);