feat add c_shim_npu-pytorch-AtomGit

已合并

feat add c_shim_npu #35292

huangyunlong创建于 5月11日从huangyunlong2022/pytorch： 2.9shim 合入到 Ascend/pytorch： v2.9.0

feat add c_shim_npu #35292

已合并

huangyunlong创建于 5月11日

从huangyunlong2022/pytorch 2.9shim合入到Ascend/pytorch v2.9.0

共 7 个文件变更+646-1

Mtest/cpp_extensions/setup.py+3-0

	@@ -20,6 +20,9 @@ ext_modules = [
20	NpuExtension(	20	NpuExtension(
21	'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],	21	'torch_test_cpp_extension.npu_from_blob', ['test_from_blob.cpp'],
22	extra_compile_args=CXX_FLAGS),	22	extra_compile_args=CXX_FLAGS),
		23	+ NpuExtension(
		24	+ 'torch_test_cpp_extension.stable_libtorch', ['test_stable_libtorch.cpp'],
		25	+ extra_compile_args=CXX_FLAGS),
23	]	26	]
24		27
25	setup(	28	setup(

Mtest/cpp_extensions/test/test_cpp_extensions_aot.py+27-0

	@@ -1,3 +1,6 @@
		1	+import ctypes
		2	+import inspect
		3	+from pathlib import Path
1	import os	4	import os
2	import stat	5	import stat
3	import pathlib	6	import pathlib
	@@ -13,6 +16,7 @@ from torch_npu.testing.testcase import TestCase, run_tests
13	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	16	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
14	from torch_npu.testing.common_utils import create_common_tensor	17	from torch_npu.testing.common_utils import create_common_tensor
15		18
		19	+
16	try:	20	try:
17	import torch_test_cpp_extension.npu as npu_extension	21	import torch_test_cpp_extension.npu as npu_extension
18	import torch_test_cpp_extension.npu_from_blob as from_blob_ext	22	import torch_test_cpp_extension.npu_from_blob as from_blob_ext
	@@ -21,6 +25,7 @@ except ImportError as e:
21	"test_cpp_extensions_aot.py cannot be invoked directly. Run "	25	"test_cpp_extensions_aot.py cannot be invoked directly. Run "
22	"`python run_cpp_test.py` instead.") from e	26	"`python run_cpp_test.py` instead.") from e
23		27
		28	+
24	class TestCppExtensionAOT(TestCase):	29	class TestCppExtensionAOT(TestCase):
25	"""Tests ahead-of-time cpp extensions	30	"""Tests ahead-of-time cpp extensions
26	"""	31	"""
	@@ -191,6 +196,7 @@ class TestCppExtensionAOT(TestCase):
191	os.remove(dump_pth)	196	os.remove(dump_pth)
192	os.remove(dump_pth + "_py_traceback")	197	os.remove(dump_pth + "_py_traceback")
193		198
		199	+
194	class TestFromBlob(TestCase):	200	class TestFromBlob(TestCase):
195	"""Tests for at_npu::native::from_blob interface"""	201	"""Tests for at_npu::native::from_blob interface"""
196		202
	@@ -221,5 +227,26 @@ class TestFromBlob(TestCase):
221	def test_from_blob_clone(self):	227	def test_from_blob_clone(self):
222	self.assertTrue(from_blob_ext.test_from_blob_clone())	228	self.assertTrue(from_blob_ext.test_from_blob_clone())
223		229
		230	+
		231	+class TestStableLibtorch(TestCase):
		232	+ @classmethod
		233	+ def setUpClass(cls):
		234	+ so_files = list(Path(inspect.getfile(npu_extension)).parent.glob("libtorch"))
		235	+ with torch._ops.dl_open_guard():
		236	+ loaded_lib = ctypes.CDLL(str(so_files[0]))
		237	+
		238	+ def test_my_abs(self):
		239	+ t = torch.rand(2).npu()
		240	+
		241	+ res = torch.ops.libtorch_agn_211.my_abs(t)
		242	+ self.assertEqual(res, torch.abs(t))
		243	+
		244	+ def test_my_cummax(self):
		245	+ t = torch.rand(2, 3, 4).npu()
		246	+
		247	+ res = torch.ops.libtorch_agn_211.my_cummax(t, 0)
		248	+ self.assertEqual(res, torch.cummax(t, 0))
		249	+
		250	+
224	if __name__ == "__main__":	251	if __name__ == "__main__":
225	run_tests()	252	run_tests()

Atest/cpp_extensions/test_stable_libtorch.cpp+43-0

@@ -0,0 +1,43 @@
	1	+#include <torch/csrc/stable/library.h>
	2	+#include <torch/csrc/stable/tensor.h>
	3	+#include "torch_npu/csrc/inductor/aoti_torch/generated/c_shim_npu.h"
	4	+
	5	+using torch::stable::Tensor;
	6	+
	7	+Tensor my_abs(Tensor self)
	8	+{
	9	+ AtenTensorHandle ret;
	10	+ aoti_torch_npu_abs(self.get(), &ret);
	11	+ return Tensor(ret);
	12	+}
	13	+
	14	+void boxed_my_abs(StableIValue* stack, uint64_t num_args, uint64_t num_outs)
	15	+{
	16	+ Tensor res = my_abs(to<Tensor>(stack[0]));
	17	+ stack[0] = from(res);
	18	+}
	19	+
	20	+std::tuple<Tensor, Tensor> my_cummax(Tensor self, int64_t dim)
	21	+{
	22	+ AtenTensorHandle ret0;
	23	+ AtenTensorHandle ret1;
	24	+ aoti_torch_npu_cummax(self.get(), dim, &ret0, &ret1);
	25	+ return std::make_tuple(Tensor(ret0), Tensor(ret1));
	26	+}
	27	+
	28	+void boxed_my_cummax(StableIValue* stack, uint64_t num_args, uint64_t num_outs)
	29	+{
	30	+ auto tuple = my_cummax(to<Tensor>(stack[0]), to<int64_t>(stack[1]));
	31	+ stack[0] = from(std::get<0>(tuple));
	32	+ stack[1] = from(std::get<1>(tuple));
	33	+}
	34	+
	35	+STABLE_TORCH_LIBRARY_FRAGMENT(libtorch_agn_211, m) {
	36	+ m.def("my_abs(Tensor self) -> Tensor");
	37	+ m.def("my_cummax(Tensor self, int dim) -> (Tensor, Tensor)");
	38	+}
	39	+
	40	+STABLE_TORCH_LIBRARY_IMPL(libtorch_agn_211, CompositeExplicitAutograd, m) {
	41	+ m.impl("my_abs", &boxed_my_abs);
	42	+ m.impl("my_cummax", &boxed_my_cummax);
	43	+}

Mtorch_npu/csrc/inductor/CMakeLists.txt+1-0

	@@ -4,6 +4,7 @@ FILE(GLOB _INDUCTOR_SRCS
4	aoti_torch/*.cpp	4	aoti_torch/*.cpp
5	aoti_package/*.cpp	5	aoti_package/*.cpp
6	dvm/*.cpp	6	dvm/*.cpp
		7	+ aoti_torch/generated/*.cpp
7	mlir/*.cpp)	8	mlir/*.cpp)
8		9
9	LIST(APPEND INDUCTOR_SRCS ${_INDUCTOR_SRCS})	10	LIST(APPEND INDUCTOR_SRCS ${_INDUCTOR_SRCS})

Atorch_npu/csrc/inductor/aoti_torch/generated/c_shim_npu.h+128-0

@@ -0,0 +1,128 @@
	1	+
	2	+
	3	+// WARNING: THIS FILE IS AUTOGENERATED BY torchnpugen. DO NOT MODIFY BY HAND.
	4	+// See torchnpugen/gen_npu_c_shim.py for details
	5	+
	6	+#pragma once
	7	+
	8	+#include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
	9	+
	10	+#ifdef __cplusplus
	11	+extern "C" {
	12	+#endif
	13	+
	14	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__adaptive_avg_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
	15	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__adaptive_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
	16	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__adaptive_avg_pool3d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
	17	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__adaptive_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
	18	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__cdist_backward(AtenTensorHandle grad, AtenTensorHandle x1, AtenTensorHandle x2, double p, AtenTensorHandle cdist, AtenTensorHandle* ret0);
	19	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__cdist_forward(AtenTensorHandle x1, AtenTensorHandle x2, double p, int64_t* compute_mode, AtenTensorHandle* ret0);
	20	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
	21	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__embedding_bag_forward_only(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
	22	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__embedding_bag_per_sample_weights_backward(AtenTensorHandle grad, AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, AtenTensorHandle offset2bag, int64_t mode, int64_t padding_idx, AtenTensorHandle* ret0);
	23	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__fft_c2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t forward, AtenTensorHandle* ret0);
	24	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__fft_r2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t onesided, AtenTensorHandle* ret0);
	25	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
	26	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	27	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu__pdist_forward(AtenTensorHandle self, double p, AtenTensorHandle* ret0);
	28	+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_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__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__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);
	32	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_abs(AtenTensorHandle self, AtenTensorHandle* ret0);
	33	+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);
	34	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_adaptive_max_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
	35	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_adaptive_max_pool3d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	36	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_adaptive_max_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
	37	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_add_Scalar(AtenTensorHandle self, double other, double alpha, AtenTensorHandle* ret0);
	38	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_add_Tensor(AtenTensorHandle self, AtenTensorHandle other, double alpha, AtenTensorHandle* ret0);
	39	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_addbmm(AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha, AtenTensorHandle* ret0);
	40	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_addmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha);
	41	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_addmv(AtenTensorHandle self, AtenTensorHandle mat, AtenTensorHandle vec, double beta, double alpha, AtenTensorHandle* ret0);
	42	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_angle(AtenTensorHandle self, AtenTensorHandle* ret0);
	43	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_avg_pool2d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
	44	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
	45	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_avg_pool3d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
	46	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
	47	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_baddbmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha);
	48	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_bernoulli__Tensor(AtenTensorHandle self, AtenTensorHandle p, AtenGeneratorHandle* generator);
	49	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_bernoulli__float(AtenTensorHandle self, double p, AtenGeneratorHandle* generator);
	50	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_bmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
	51	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_bucketize_Tensor(AtenTensorHandle self, AtenTensorHandle boundaries, int32_t out_int32, int32_t right, AtenTensorHandle* ret0);
	52	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_cat(const AtenTensorHandle* tensors, int64_t tensors_len_, int64_t dim, AtenTensorHandle* ret0);
	53	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_cholesky_solve(AtenTensorHandle self, AtenTensorHandle input2, int32_t upper, AtenTensorHandle* ret0);
	54	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_convolution(AtenTensorHandle input, AtenTensorHandle weight, AtenTensorHandle* bias, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, AtenTensorHandle* ret0);
	55	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_convolution_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle weight, const int64_t** bias_sizes, int64_t bias_sizes_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
	56	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_cummax(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	57	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_cummin(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	58	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_cumsum(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
	59	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_exponential(AtenTensorHandle self, double lambd, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
	60	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_fill__Scalar(AtenTensorHandle self, double value);
	61	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_grid_sampler_2d_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle grid, int64_t interpolation_mode, int64_t padding_mode, int32_t align_corners, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	62	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	63	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_histc(AtenTensorHandle self, int64_t bins, double min, double max, AtenTensorHandle* ret0);
	64	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_index_Tensor(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle* ret0);
	65	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_index_put(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle values, int32_t accumulate, AtenTensorHandle* ret0);
	66	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_kthvalue(AtenTensorHandle self, int64_t k, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	67	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_logcumsumexp(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
	68	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_masked_scatter(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle source, AtenTensorHandle* ret0);
	69	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_masked_scatter_backward(AtenTensorHandle grad_output, AtenTensorHandle mask, const int64_t* sizes, int64_t sizes_len_, AtenTensorHandle* ret0);
	70	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_masked_select(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle* ret0);
	71	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_max_pool2d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	72	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_max_pool2d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
	73	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_max_pool3d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	74	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_max_pool3d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
	75	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_max_unpool2d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
	76	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_max_unpool3d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
	77	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_median(AtenTensorHandle self, AtenTensorHandle* ret0);
	78	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
	79	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_mul_Scalar(AtenTensorHandle self, double other, AtenTensorHandle* ret0);
	80	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_mul_Tensor(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
	81	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_nanmedian(AtenTensorHandle self, AtenTensorHandle* ret0);
	82	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
	83	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_native_dropout(AtenTensorHandle input, double p, int32_t* train, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	84	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_nonzero(AtenTensorHandle self, AtenTensorHandle* ret0);
	85	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
	86	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
	87	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0);
	88	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0);
	89	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_pow_Scalar(double self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
	90	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_pow_Tensor_Scalar(AtenTensorHandle self, double exponent, AtenTensorHandle* ret0);
	91	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
	92	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	93	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	94	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	95	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	96	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	97	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
	98	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	99	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	100	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
	101	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
	102	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_replication_pad2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
	103	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0);
	104	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_resize_(AtenTensorHandle self, const int64_t* size, int64_t size_len_, int32_t* memory_format);
	105	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format);
	106	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_scatter_src_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src);
	107	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_scatter_value_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, double value);
	108	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_searchsorted_Scalar(AtenTensorHandle sorted_sequence, double self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
	109	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_searchsorted_Tensor(AtenTensorHandle sorted_sequence, AtenTensorHandle self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
	110	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_set__source_Tensor(AtenTensorHandle self, AtenTensorHandle source);
	111	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_slice_Tensor(AtenTensorHandle self, int64_t dim, int64_t* start, int64_t* end, int64_t step, AtenTensorHandle* ret0);
	112	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_soft_margin_loss_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle target, int64_t reduction, AtenTensorHandle* ret0);
	113	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_sort(AtenTensorHandle self, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	114	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_sort_stable(AtenTensorHandle self, int32_t* stable, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	115	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_squeeze_dim(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
	116	+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);
	117	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_topk(AtenTensorHandle self, int64_t k, int64_t dim, int32_t largest, int32_t sorted, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
	118	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_uniform(AtenTensorHandle self, double from, double to, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
	119	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_upsample_bicubic2d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
	120	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_upsample_linear1d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales, AtenTensorHandle* ret0);
	121	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_upsample_trilinear3d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_d, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
	122	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_view_dtype(AtenTensorHandle self, int32_t dtype, AtenTensorHandle* ret0);
	123	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_view_as_complex(AtenTensorHandle self, AtenTensorHandle* ret0);
	124	+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_npu_view_as_real(AtenTensorHandle self, AtenTensorHandle* ret0);
	125	+
	126	+#ifdef __cplusplus
	127	+} // extern "C"
	128	+#endif

Mtorchnpugen/gen_backend_stubs.py+138-1

	@@ -61,6 +61,7 @@ from torchnpugen.gen_functionalization_type import (
61	gen_functionalization_definition,	61	gen_functionalization_definition,
62	gen_functionalization_registration,	62	gen_functionalization_registration,
63	)	63	)
		64	+from torchnpugen.gen_npu_c_shim import gen_npu_c_shim_files
64	from torchnpugen.utils import (	65	from torchnpugen.utils import (
65	add_header_to_template_file,	66	add_header_to_template_file,
66	DEVICE_NOCHECK_SET,	67	DEVICE_NOCHECK_SET,
	@@ -539,6 +540,12 @@ def main() -> None:
539	default=None,	540	default=None,
540	help="path to the source yaml file containing kernel definitions in op_plugin",	541	help="path to the source yaml file containing kernel definitions in op_plugin",
541	)	542	)
		543	+ parser.add_argument(
		544	+ "--update_aoti_c_shim",
		545	+ action="store_true",
		546	+ help="Update AOTInductor C shim after adding an entry to inductor_fallback_ops in torchgen/aoti/fallback_ops.py. "
		547	+ "WARNING: Do not use this unless you are sure what you are doing!!!",
		548	+ )
542	options = parser.parse_args()	549	options = parser.parse_args()
543		550
544	run(	551	run(
	@@ -548,6 +555,7 @@ def main() -> None:
548	options.impl_path,	555	options.impl_path,
549	options.op_plugin_impl_path,	556	options.op_plugin_impl_path,
550	options.op_plugin_yaml_path,	557	options.op_plugin_yaml_path,
		558	+ options.update_aoti_c_shim,
551	)	559	)
552		560
553		561
	@@ -637,7 +645,20 @@ $dispatch_registrations_body
637	"static_init_dispatch_registrations": static_init_dispatch_registrations,	645	"static_init_dispatch_registrations": static_init_dispatch_registrations,
638	"deferred_dispatch_registrations": "",	646	"deferred_dispatch_registrations": "",
639	"dispatch_namespace": dispatch_key.lower(),	647	"dispatch_namespace": dispatch_key.lower(),
640	- "dispatch_namespaced_definitions": native_function_registrations,	648	+ "dispatch_namespaced_definitions": native_function_registrations if native_function_registrations else list(
		649	+ concatMap(
		650	+ register_dispatch_key_func(
		651	+ backend_index,
		652	+ Target.NAMESPACED_DEFINITION,
		653	+ selector,
		654	+ rocm=False,
		655	+ symint=True,
		656	+ class_method_name=f'{class_name}',
		657	+ skip_dispatcher_op_registration=False,
		658	+ ),
		659	+ grouped_native_functions,
		660	+ )
		661	+ ),
641	"dispatch_anonymous_definitions": list(	662	"dispatch_anonymous_definitions": list(
642	concatMap(	663	concatMap(
643	register_dispatch_key_func(	664	register_dispatch_key_func(
	@@ -948,6 +969,102 @@ def gen_target_registration(
948	)	969	)
949		970
950		971
		972	+def gen_per_operator_headers(
		973	+ fm: FileManager,
		974	+ ops_fm: FileManager,
		975	+ native_functions: Sequence[NativeFunction],
		976	+ grouped_native_functions: Sequence[NativeFunction \| NativeFunctionsGroup],
		977	+ backend_indices: dict[DispatchKey, BackendIndex],
		978	+ dispatch_keys: Sequence[DispatchKey],
		979	+ selector: "SelectiveBuilder",
		980	+):
		981	+ """
		982	+ Generate per-operator dispatch header files (*_npu_dispatch.h) for NPU.
		983	+
		984	+ This mirrors upstream's gen_per_operator_headers which generates
		985	+ _cuda_dispatch.h and _cpu_dispatch.h in ATen/ops/.
		986	+
		987	+ Uses dest.RegisterDispatchKey with Target.NAMESPACED_DECLARATION to
		988	+ generate TORCH_API function declarations in at::npu namespace,
		989	+ exactly matching upstream's dispatch header format.
		990	+
		991	+ Also generates NPUFunctions.h/NPUFunctions_inl.h which includes all per-operator
		992	+ dispatch headers, mirroring upstream's CUDAFunctions.h/CUDAFunctions_inl.h.
		993	+ """
		994	+
		995	+ functions_by_root_name: dict[str, list[NativeFunction]] = defaultdict(list)
		996	+ for fn in native_functions:
		997	+ functions_by_root_name[fn.root_name].append(fn)
		998	+
		999	+ grouped_functions_by_root_name: dict[
		1000	+ str, list[NativeFunction \| NativeFunctionsGroup]
		1001	+ ] = defaultdict(list)
		1002	+ for group in grouped_native_functions:
		1003	+ name = group.root_name
		1004	+ grouped_functions_by_root_name[name].append(group)
		1005	+
		1006	+ for dispatch_key in dispatch_keys:
		1007	+ if dispatch_key not in backend_indices:
		1008	+ continue
		1009	+
		1010	+ dispatch_namespace = dispatch_key.lower()
		1011	+ dispatch_names = []
		1012	+
		1013	+ for name, functions in functions_by_root_name.items():
		1014	+ grouped_functions = grouped_functions_by_root_name.get(name, [])
		1015	+ declarations = list(
		1016	+ concatMap(
		1017	+ dest.RegisterDispatchKey(
		1018	+ backend_indices[dispatch_key],
		1019	+ Target.NAMESPACED_DECLARATION,
		1020	+ selector,
		1021	+ rocm=False,
		1022	+ symint=True,
		1023	+ class_method_name=None,
		1024	+ skip_dispatcher_op_registration=False,
		1025	+ ),
		1026	+ grouped_functions,
		1027	+ )
		1028	+ )
		1029	+
		1030	+ if len(declarations) == 0:
		1031	+ continue
		1032	+
		1033	+ dispatch_names.append(name)
		1034	+
		1035	+ ops_fm.write_with_template(
		1036	+ f"{name}_{dispatch_namespace}_dispatch.h",
		1037	+ "DispatchKeyFunction.h",
		1038	+ lambda: {
		1039	+ "dispatch_namespace": dispatch_namespace,
		1040	+ "dispatch_namespaced_declarations": declarations,
		1041	+ },
		1042	+ )
		1043	+
		1044	+ inl_headers = f"#include <torch_npu/csrc/aten/{dispatch_key}Functions_inl.h>"
		1045	+
		1046	+ fm.write_with_template(
		1047	+ f"{dispatch_key}Functions.h",
		1048	+ "DispatchKeyFunctions.h",
		1049	+ lambda: {
		1050	+ "dispatch_key": str(dispatch_key),
		1051	+ "inline_headers": inl_headers,
		1052	+ },
		1053	+ )
		1054	+ fm.write_with_template(
		1055	+ f"{dispatch_key}Functions_inl.h",
		1056	+ "DispatchKeyFunctions_inl.h",
		1057	+ lambda: {
		1058	+ "dispatch_namespace": dispatch_namespace,
		1059	+ "DispatchKeyFunctions_inl_includes": [
		1060	+ f"#include <torch_npu/csrc/aten/ops/{name}_{dispatch_namespace}_dispatch.h>"
		1061	+ for name in sorted(dispatch_names)
		1062	+ ],
		1063	+ "dispatch_namespaced_declarations": [],
		1064	+ },
		1065	+ )
		1066	+
		1067	+
951	def run(	1068	def run(
952	source_yaml: str,	1069	source_yaml: str,
953	output_dir: str,	1070	output_dir: str,
	@@ -955,6 +1072,7 @@ def run(
955	impl_path: str \| None,	1072	impl_path: str \| None,
956	op_plugin_impl_path: str \| None,	1073	op_plugin_impl_path: str \| None,
957	op_plugin_yaml_path: str \| None,	1074	op_plugin_yaml_path: str \| None,
		1075	+ update_aoti_c_shim: bool,
958	) -> None:	1076	) -> None:
959	rename_privateuse1_dispatch_key()	1077	rename_privateuse1_dispatch_key()
960	torchgen_path = get_torchgen_dir()	1078	torchgen_path = get_torchgen_dir()
	@@ -967,6 +1085,10 @@ def run(
967	)	1085	)
968		1086
969	fm = make_file_manager(output_dir)	1087	fm = make_file_manager(output_dir)
		1088	+ ops_output_dir = os.path.join(output_dir, "ops")
		1089	+ if not os.path.exists(ops_output_dir):
		1090	+ os.makedirs(ops_output_dir, exist_ok=True)
		1091	+ ops_fm = make_file_manager(ops_output_dir)
970	merge_custom_yaml(source_yaml, op_plugin_yaml_path)	1092	merge_custom_yaml(source_yaml, op_plugin_yaml_path)
971	source_yaml = gen_custom_yaml_path(source_yaml)	1093	source_yaml = gen_custom_yaml_path(source_yaml)
972	tags_yaml_path = os.path.join(torchgen_path, "packaged/ATen/native/tags.yaml")	1094	tags_yaml_path = os.path.join(torchgen_path, "packaged/ATen/native/tags.yaml")
	@@ -1015,6 +1137,8 @@ def run(
1015	None,	1137	None,
1016	)	1138	)
1017		1139
		1140	+ gen_per_operator_headers(fm, ops_fm, native_functions, grouped_native_functions, backend_indices, [backend_dispatch_key, autograd_dispatch_key], selector)
		1141	+
1018	for dispatch_key in [backend_dispatch_key, autograd_dispatch_key]:	1142	for dispatch_key in [backend_dispatch_key, autograd_dispatch_key]:
1019	if not dispatch_key:	1143	if not dispatch_key:
1020	continue	1144	continue
	@@ -1117,6 +1241,19 @@ def run(
1117	native_functions,	1241	native_functions,
1118	)	1242	)
1119		1243
		1244	+ aoti_output_dir = os.path.join(output_dir, "../inductor/aoti_torch/generated")
		1245	+ if not os.path.exists(aoti_output_dir):
		1246	+ os.makedirs(aoti_output_dir, exist_ok=True)
		1247	+ aoti_fm = make_file_manager(aoti_output_dir)
		1248	+ structured_native_functions = [
		1249	+ g for g in grouped_native_functions if isinstance(g, NativeFunctionsGroup)
		1250	+ ]
		1251	+ gen_npu_c_shim_files(
		1252	+ aoti_fm, native_functions, backend_indices,
		1253	+ [backend_dispatch_key, autograd_dispatch_key],
		1254	+ structured_native_functions, update_aoti_c_shim,
		1255	+ )
		1256	+
1120		1257
1121	def apply_torchgen_patch():	1258	def apply_torchgen_patch():
1122	dest.RegisterDispatchKey.gen_unstructured = gen_unstructured	1259	dest.RegisterDispatchKey.gen_unstructured = gen_unstructured

Atorchnpugen/gen_npu_c_shim.py+306-0

@@ -0,0 +1,306 @@
	1	+"""
	2	+NPU C Shim Code Generator for AOTInductor.
	3	+
	4	+Generates c_shim_npu.h and c_shim_npu.cpp for the NPU device,
	5	+following the same pattern as PyTorch's torchgen/gen_aoti_c_shim.py
	6	+which generates c_shim_cuda.h/cpp and c_shim_cpu.h/cpp.
	7	+
	8	+Integrated into the torchnpugen codegen pipeline.
	9	+
	10	+Static dispatch:
	11	+ Following upstream pattern, we use at::npu::* instead of at::* to bypass the dispatcher
	12	+ and directly call NPU backend kernels for better performance.
	13	+"""
	14	+
	15	+from __future__ import annotations
	16	+
	17	+import difflib
	18	+import os
	19	+import textwrap
	20	+from collections.abc import Sequence
	21	+from dataclasses import dataclass
	22	+
	23	+from torchgen.aoti.fallback_ops import inductor_fallback_ops
	24	+from torchgen.context import method_with_native_function
	25	+from torchgen.gen import FileManager
	26	+from torchgen.gen_aoti_c_shim import (
	27	+ gen_declaration_and_definition,
	28	+ gen_static_dispatch_backend_call,
	29	+ get_fallback_op_name,
	30	+)
	31	+from torchgen.model import (
	32	+ BackendIndex,
	33	+ DispatchKey,
	34	+ NativeFunction,
	35	+ NativeFunctionsGroup,
	36	+ OperatorName,
	37	+)
	38	+from torchgen.utils import mapMaybe
	39	+
	40	+
	41	+NPU_DEVICE = "npu"
	42	+
	43	+NPU_DISPATCH_KEYS: Sequence[DispatchKey] = ()
	44	+
	45	+
	46	+def _get_backend_index_for_npu(
	47	+ func: NativeFunction,
	48	+ backend_indices: dict[DispatchKey, BackendIndex],
	49	+ structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup] = {},
	50	+) -> BackendIndex \| None:
	51	+ for dk in NPU_DISPATCH_KEYS:
	52	+ if dk in backend_indices:
	53	+ if backend_indices[dk].has_kernel(func) or (
	54	+ func.structured_delegate is not None
	55	+ and func.structured_delegate in structured_func_group_dict
	56	+ and backend_indices[dk].has_kernel(
	57	+ structured_func_group_dict[func.structured_delegate]
	58	+ )
	59	+ ):
	60	+ return backend_indices[dk]
	61	+ if DispatchKey.CompositeExplicitAutograd in backend_indices:
	62	+ if backend_indices[DispatchKey.CompositeExplicitAutograd].has_kernel(func):
	63	+ return backend_indices[DispatchKey.CompositeExplicitAutograd]
	64	+ if DispatchKey.CompositeExplicitAutogradNonFunctional in backend_indices:
	65	+ if backend_indices[
	66	+ DispatchKey.CompositeExplicitAutogradNonFunctional
	67	+ ].has_kernel(func):
	68	+ return backend_indices[DispatchKey.CompositeExplicitAutogradNonFunctional]
	69	+ if DispatchKey.CompositeImplicitAutograd in backend_indices:
	70	+ if backend_indices[DispatchKey.CompositeImplicitAutograd].has_kernel(func):
	71	+ return backend_indices[DispatchKey.CompositeImplicitAutograd]
	72	+ return None
	73	+
	74	+
	75	+def _gen_npu_c_shim(
	76	+ func: NativeFunction,
	77	+ version_info: dict[str, list[str]],
	78	+ backend_indices: dict[DispatchKey, BackendIndex],
	79	+ structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup],
	80	+ header: bool,
	81	+) -> str \| None:
	82	+ backend_index = _get_backend_index_for_npu(
	83	+ func, backend_indices, structured_func_group_dict
	84	+ )
	85	+ if backend_index is None:
	86	+ return None
	87	+
	88	+ schema = func.func
	89	+ device = NPU_DEVICE
	90	+ backend_call = gen_static_dispatch_backend_call(func, backend_index)
	91	+
	92	+ try:
	93	+ if header:
	94	+ declaration, _ = gen_declaration_and_definition(
	95	+ schema, device, backend_call, version_info
	96	+ )
	97	+ return declaration
	98	+ else:
	99	+ _, definition = gen_declaration_and_definition(
	100	+ schema, device, backend_call, version_info
	101	+ )
	102	+ return definition
	103	+ except NotImplementedError:
	104	+ return None
	105	+
	106	+
	107	+def _get_dispatch_header_path(
	108	+ func: NativeFunction,
	109	+ backend_index: BackendIndex \| None,
	110	+) -> str \| None:
	111	+ if backend_index is None:
	112	+ return None
	113	+ dispatch_ns = backend_index.dispatch_key.lower()
	114	+ return (
	115	+ f"#include <torch_npu/csrc/aten/ops/{func.root_name}_{dispatch_ns}_dispatch.h>"
	116	+ )
	117	+
	118	+
	119	+@dataclass(frozen=True)
	120	+class NPUShimGenerator:
	121	+ npu_fallback_ops: dict[str, dict[str, list[str]]]
	122	+ backend_indices: dict[DispatchKey, BackendIndex]
	123	+ structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup]
	124	+ header: bool
	125	+
	126	+ @method_with_native_function
	127	+ def __call__(self, func: NativeFunction) -> str \| None:
	128	+ version_info = self.npu_fallback_ops[get_fallback_op_name(func)]
	129	+ return _gen_npu_c_shim(
	130	+ func,
	131	+ version_info,
	132	+ self.backend_indices,
	133	+ self.structured_func_group_dict,
	134	+ self.header,
	135	+ )
	136	+
	137	+
	138	+def gen_npu_c_shim(
	139	+ native_functions: Sequence[NativeFunction],
	140	+ npu_fallback_ops: dict[str, dict[str, list[str]]],
	141	+ backend_indices: dict[DispatchKey, BackendIndex],
	142	+ structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup],
	143	+ header: bool,
	144	+ includes: str = "",
	145	+) -> str:
	146	+ body = "\n".join(
	147	+ list(
	148	+ mapMaybe(
	149	+ NPUShimGenerator(
	150	+ npu_fallback_ops,
	151	+ backend_indices,
	152	+ structured_func_group_dict,
	153	+ header,
	154	+ ),
	155	+ native_functions,
	156	+ )
	157	+ )
	158	+ )
	159	+
	160	+ warning = """
	161	+
	162	+// WARNING: THIS FILE IS AUTOGENERATED BY torchnpugen. DO NOT MODIFY BY HAND.
	163	+// See torchnpugen/gen_npu_c_shim.py for details"""
	164	+
	165	+ if header:
	166	+ return (
	167	+ warning
	168	+ + textwrap.dedent("""
	169	+
	170	+ #pragma once
	171	+
	172	+ #include <torch_npu/csrc/inductor/aoti_torch/c/shim.h>
	173	+
	174	+ #ifdef __cplusplus
	175	+ extern "C" {
	176	+ #endif
	177	+
	178	+ """)
	179	+ + body
	180	+ + textwrap.dedent("""
	181	+
	182	+ #ifdef __cplusplus
	183	+ } // extern "C"
	184	+ #endif
	185	+ """)
	186	+ )
	187	+ else:
	188	+ return (
	189	+ warning
	190	+ + textwrap.dedent(f"""
	191	+
	192	+ #include <torch_npu/csrc/inductor/aoti_torch/generated/c_shim_{NPU_DEVICE}.h>
	193	+ #include <torch_npu/csrc/inductor/aoti_torch/utils.h>
	194	+
	195	+ #ifndef AT_PER_OPERATOR_HEADERS
	196	+ #include <torch_npu/csrc/aten/NPUFunctions.h>
	197	+ #include <ATen/CompositeExplicitAutogradFunctions.h>
	198	+ #include <ATen/CompositeExplicitAutogradNonFunctionalFunctions.h>
	199	+ #include <ATen/CompositeImplicitAutogradFunctions.h>
	200	+ #else
	201	+ """)
	202	+ + includes
	203	+ + textwrap.dedent("""
	204	+ #endif // AT_PER_OPERATOR_HEADERS
	205	+
	206	+ using namespace torch::aot_inductor;
	207	+
	208	+ """)
	209	+ + body
	210	+ )
	211	+
	212	+
	213	+def gen_npu_c_shim_files(
	214	+ aoti_fm: FileManager,
	215	+ native_functions: Sequence[NativeFunction],
	216	+ backend_indices: dict[DispatchKey, BackendIndex],
	217	+ dispatch_keys: Sequence[DispatchKey],
	218	+ structured_native_functions: Sequence[NativeFunctionsGroup],
	219	+ update_aoti_c_shim: bool,
	220	+) -> None:
	221	+ global NPU_DISPATCH_KEYS
	222	+ NPU_DISPATCH_KEYS = dispatch_keys
	223	+
	224	+ structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup] = {}
	225	+ for func_group in structured_native_functions:
	226	+ for func in func_group.functions():
	227	+ if func.structured_delegate is not None:
	228	+ structured_func_group_dict[func.structured_delegate] = func_group
	229	+ break
	230	+
	231	+ fallback_ops_dict = inductor_fallback_ops
	232	+ fallbacks = {}
	233	+ for func in native_functions:
	234	+ op_name = get_fallback_op_name(func)
	235	+ if op_name in fallback_ops_dict:
	236	+ fallbacks[op_name] = func
	237	+ fallback_native_functions = tuple(value for _, value in sorted(fallbacks.items()))
	238	+
	239	+ def headers_for_npu() -> str:
	240	+ headers = []
	241	+ for func in fallback_native_functions:
	242	+ backend_index = _get_backend_index_for_npu(
	243	+ func, backend_indices, structured_func_group_dict
	244	+ )
	245	+ header = _get_dispatch_header_path(func, backend_index)
	246	+ if header is not None:
	247	+ headers.append(header)
	248	+ return "\n".join(sorted(set(headers)))
	249	+
	250	+ header_content = gen_npu_c_shim(
	251	+ fallback_native_functions,
	252	+ fallback_ops_dict,
	253	+ backend_indices,
	254	+ structured_func_group_dict,
	255	+ header=True,
	256	+ )
	257	+ cpp_content = gen_npu_c_shim(
	258	+ fallback_native_functions,
	259	+ fallback_ops_dict,
	260	+ backend_indices,
	261	+ structured_func_group_dict,
	262	+ header=False,
	263	+ includes=headers_for_npu(),
	264	+ )
	265	+
	266	+ header_filename = f"c_shim_{NPU_DEVICE}.h"
	267	+ cpp_filename = f"c_shim_{NPU_DEVICE}.cpp"
	268	+
	269	+ if update_aoti_c_shim:
	270	+ aoti_fm.write(header_filename, lambda: header_content)
	271	+ else:
	272	+ try:
	273	+ with open(os.path.join(aoti_fm.install_dir, header_filename)) as old_file:
	274	+ old_header = old_file.read()
	275	+ if old_header != header_content:
	276	+ diff = "\n".join(
	277	+ difflib.unified_diff(
	278	+ old_header.splitlines(),
	279	+ header_content.splitlines(),
	280	+ fromfile="expected",
	281	+ tofile="actual",
	282	+ lineterm="",
	283	+ )
	284	+ )
	285	+ raise RuntimeError(f"""
	286	+The generated AOTInductor C shim header files have unexpectedly changed. This
	287	+indicates an AOTInductor fallback operator ABI backward compatibility breakage!!!
	288	+
	289	+1. You added a fallback op to the inductor_fallback_ops list in torchgen/aoti/fallback_ops.py.
	290	+If that's the case, run codegen with --update-aoti-c-shim to add a new entry to
	291	+existing C shim header files.
	292	+
	293	+2. You added a new default argument to an existing fallback op. This is clearly a BC breaking
	294	+change in the AOTInductor land. You need to annotate the new default argument in
	295	+torchgen/aoti/fallback_ops.py, and then run codegen with --update-aoti-c-shim to
	296	+update the C shim header files by creating different versions of the fallback op.
	297	+
	298	+{diff}
	299	+ """)
	300	+ except FileNotFoundError:
	301	+ print(f"{os.path.join(aoti_fm.install_dir, header_filename)} not found")
	302	+
	303	+ aoti_fm.write(cpp_filename, lambda: cpp_content)
	304	+
	305	+ print(f"[torchnpugen] Generated {header_filename}")
	306	+ print(f"[torchnpugen] Generated {cpp_filename}")