Star863
1359
代码
代码
Issues1725
Pull Requests2246
项目讨论
Wiki
分析
Star863
1359
已合并
[Inductor] remove inductor static mode and interface.cpp in inductor #39997
zhucehw创建于 7月3日
[Inductor] remove inductor static mode and interface.cpp in inductor #39997
已合并
zhucehw创建于 7月3日
从zhucehw/pytorchv2.10.0-dev合入到Ascend/pytorchv2.10.0
共 6 个文件变更+49-365
Mtest/_inductor/testutils.py+1-1
@@ -14,7 +14,7 @@ class TestUtils(TestCase):
14 (32, 32, 1024, 1024)] # 128*128*4096*2048 is too big(512G)14 (32, 32, 1024, 1024)] # 128*128*4096*2048 is too big(512G)
15 _pointwise_test_shapes = _pointwise_test_shape2d + _pointwise_test_shape3d + _pointwise_test_shape4d15 _pointwise_test_shapes = _pointwise_test_shape2d + _pointwise_test_shape3d + _pointwise_test_shape4d
16 16 
17- _pointwise_demo_shapes = [(1024, 32), (8, 16, 256, 32)]17+ _pointwise_demo_shapes = [(1024, 32), (8, 16, 32)]
18 _reduction_extest_shape4d = [(8, 8, 8, 16384), (8, 8, 16384, 8), (8, 16384, 8, 8), (16384, 8, 8, 8)]18 _reduction_extest_shape4d = [(8, 8, 8, 16384), (8, 8, 16384, 8), (8, 16384, 8, 8), (16384, 8, 8, 8)]
19 _reduction_extest_dim4d = [-1, -2, 1, 0]19 _reduction_extest_dim4d = [-1, -2, 1, 0]
20 _reduction_extest_SDbinding = list(zip(_reduction_extest_shape4d, _reduction_extest_dim4d))20 _reduction_extest_SDbinding = list(zip(_reduction_extest_shape4d, _reduction_extest_dim4d))
Dtorch_npu/_inductor/codegen/aoti_runtime/interface.cpp+0-294
@@ -1,294 +0,0 @@
1-// Definition of NPU AOTI runtime interface functions
2- 
3-#include <torch/csrc/inductor/aoti_runtime/interface.h>
4-#include <torch_npu/csrc/inductor/aoti_runtime/model_container.h>
5- 
6-#include <iostream>
7-#include <sstream>
8-#include <stdexcept>
9-#include <vector>
10- 
11-#define CONVERT_EXCEPTION_TO_ERROR_CODE(...) \
12- try { \
13- __VA_ARGS__ \
14- } catch (const std::exception& e) { \
15- std::cerr << "Error: " << e.what() << std::endl; \
16- return AOTI_RUNTIME_FAILURE; \
17- } catch (...) { \
18- std::cerr << "Unknown exception occurred." << std::endl; \
19- return AOTI_RUNTIME_FAILURE; \
20- } \
21- return AOTI_RUNTIME_SUCCESS;
22- 
23-#define AOTI_VECTOR_SIZE_CHECK(actual_size, expected_size, name) \
24- do { \
25- AOTI_RUNTIME_CHECK(actual_size == expected_size, "expected " + std::string(name) + " vector size to be " + \
26- std::to_string(expected_size) + ", but got " + \
27- std::to_string(actual_size)); \
28- } while (0)
29- 
30-// AOTInductor uses at::addmm_out, which doesn't supports
31-// arguments that requires gradient. For this reason, we
32-// enforce no_grad context for run APIs.
33-//
34-// A RAII, thread local (!) guard that enables or disables grad mode upon
35-// construction, and sets it back to the original value upon destruction.
36-struct AOTINoGradGuard {
37- AOTINoGradGuard() : prev_mode(aoti_torch_grad_mode_is_enabled()) { aoti_torch_grad_mode_set_enabled(false); }
38- ~AOTINoGradGuard() { aoti_torch_grad_mode_set_enabled(prev_mode); }
39- bool prev_mode;
40-};
41- 
42-extern "C" {
43-AOTIRuntimeError AOTInductorModelContainerCreate(AOTInductorModelContainerHandle* container_handle, size_t num_models,
44- bool is_cpu, const char* cubin_dir)
45-{
46- return AOTInductorModelContainerCreateWithDevice(container_handle, num_models, is_cpu ? "cpu" : "npu", cubin_dir);
47-}
48- 
49-AOTIRuntimeError AOTInductorModelContainerCreateWithDevice(AOTInductorModelContainerHandle* container_handle,
50- size_t num_models, const char* device_str,
51- const char* cubin_dir)
52-{
53- if (num_models == 0) {
54- std::cerr << "Error: num_models must be positive, but got 0" << std::endl;
55- return AOTI_RUNTIME_FAILURE;
56- }
57- 
58- CONVERT_EXCEPTION_TO_ERROR_CODE({
59- std::optional<std::string> cubin_dir_opt;
60- if (cubin_dir != nullptr) {
61- cubin_dir_opt.emplace(cubin_dir);
62- }
63- auto* container =
64- new torch::aot_inductor::AOTInductorModelContainer(num_models, std::string(device_str), cubin_dir_opt);
65- *container_handle = reinterpret_cast<AOTInductorModelContainerHandle>(container);
66- })
67-}
68- 
69-AOTIRuntimeError AOTInductorModelContainerDelete(AOTInductorModelContainerHandle container_handle)
70-{
71- CONVERT_EXCEPTION_TO_ERROR_CODE({
72- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
73- delete container;
74- });
75-}
76- 
77-AOTIRuntimeError AOTInductorModelContainerRun(AOTInductorModelContainerHandle container_handle,
78- AtenTensorHandle* input_handles, size_t num_inputs,
79- AtenTensorHandle* output_handles, size_t num_outputs,
80- AOTInductorStreamHandle stream_handle,
81- AOTIProxyExecutorHandle proxy_executor_handle)
82-{
83- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
84- AOTI_VECTOR_SIZE_CHECK(num_inputs, container->num_inputs(), "inputs");
85- AOTI_VECTOR_SIZE_CHECK(num_outputs, container->num_outputs(), "outputs");
86- 
87- auto stream = reinterpret_cast<torch::aot_inductor::DeviceStreamType>(stream_handle);
88- CONVERT_EXCEPTION_TO_ERROR_CODE({
89- AOTINoGradGuard guard;
90- container->run(input_handles, output_handles, stream, proxy_executor_handle);
91- })
92-}
93- 
94-AOTIRuntimeError AOTInductorModelContainerRunSingleThreaded(AOTInductorModelContainerHandle container_handle,
95- AtenTensorHandle* input_handles, size_t num_inputs,
96- AtenTensorHandle* output_handles, size_t num_outputs,
97- AOTInductorStreamHandle stream_handle,
98- AOTIProxyExecutorHandle proxy_executor_handle)
99-{
100- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
101- AOTI_VECTOR_SIZE_CHECK(num_inputs, container->num_inputs(), "inputs");
102- AOTI_VECTOR_SIZE_CHECK(num_outputs, container->num_outputs(), "outputs");
103- 
104- auto stream = reinterpret_cast<torch::aot_inductor::DeviceStreamType>(stream_handle);
105- CONVERT_EXCEPTION_TO_ERROR_CODE({
106- AOTINoGradGuard guard;
107- container->run_single_threaded(input_handles, output_handles, stream, proxy_executor_handle);
108- })
109-}
110- 
111-AOTIRuntimeError AOTInductorModelContainerGetNumConstants(AOTInductorModelContainerHandle container_handle,
112- size_t* num_constants)
113-{
114- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
115- CONVERT_EXCEPTION_TO_ERROR_CODE({ *num_constants = container->num_constants(); })
116-}
117- 
118-AOTIRuntimeError AOTInductorModelContainerGetConstantName(AOTInductorModelContainerHandle container_handle, size_t idx,
119- const char** name)
120-{
121- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
122- CONVERT_EXCEPTION_TO_ERROR_CODE({ *name = container->constant_name(idx); })
123-}
124- 
125-AOTIRuntimeError AOTInductorModelContainerGetConstantOriginalFQN(AOTInductorModelContainerHandle container_handle,
126- size_t idx, const char** original_fqn)
127-{
128- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
129- CONVERT_EXCEPTION_TO_ERROR_CODE({ *original_fqn = container->constant_original_fqn(idx); })
130-}
131- 
132-AOTIRuntimeError AOTInductorModelContainerGetConstantFromFolded(AOTInductorModelContainerHandle container_handle,
133- size_t idx, bool* from_folded)
134-{
135- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
136- CONVERT_EXCEPTION_TO_ERROR_CODE({ *from_folded = container->constant_from_folded(idx); })
137-}
138- 
139-AOTIRuntimeError AOTInductorModelContainerGetConstantType(AOTInductorModelContainerHandle container_handle, size_t idx,
140- int32_t* type)
141-{
142- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
143- CONVERT_EXCEPTION_TO_ERROR_CODE({ *type = container->constant_type(idx); })
144-}
145- 
146-AOTIRuntimeError AOTInductorModelContainerGetConstantDtype(AOTInductorModelContainerHandle container_handle, size_t idx,
147- int32_t* dtype)
148-{
149- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
150- CONVERT_EXCEPTION_TO_ERROR_CODE({ *dtype = container->constant_dtype(idx); })
151-}
152- 
153-AOTIRuntimeError AOTInductorModelContainerUpdateConstantBuffer(AOTInductorModelContainerHandle container_handle,
154- AOTInductorConstantMapHandle constant_map_handle,
155- bool use_inactive, bool validate_full_update)
156-{
157- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
158- auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
159- CONVERT_EXCEPTION_TO_ERROR_CODE(
160- { container->update_constant_buffer(*input_map, use_inactive, validate_full_update); })
161-}
162- 
163-AOTIRuntimeError AOTInductorModelContainerUpdateInactiveConstantBuffer(AOTInductorModelContainerHandle container_handle,
164- AOTInductorConstantMapHandle constant_map_handle)
165-{
166- return AOTInductorModelContainerUpdateConstantBuffer(container_handle, constant_map_handle, true, true);
167-}
168- 
169-AOTIRuntimeError AOTInductorModelContainerRunConstantFolding(AOTInductorModelContainerHandle container_handle,
170- bool use_inactive, AOTInductorStreamHandle stream_handle,
171- AOTIProxyExecutorHandle proxy_executor_handle)
172-{
173- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
174- auto stream = reinterpret_cast<torch::aot_inductor::DeviceStreamType>(stream_handle);
175- CONVERT_EXCEPTION_TO_ERROR_CODE({
176- AOTINoGradGuard guard;
177- container->run_const_fold(use_inactive, stream, proxy_executor_handle);
178- })
179-}
180- 
181-AOTIRuntimeError AOTInductorModelContainerSwapConstantBuffer(AOTInductorModelContainerHandle container_handle)
182-{
183- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
184- CONVERT_EXCEPTION_TO_ERROR_CODE({ container->swap_constant_buffer(); })
185-}
186- 
187-AOTIRuntimeError AOTInductorModelContainerGetNumInputs(AOTInductorModelContainerHandle container_handle,
188- size_t* ret_num_inputs)
189-{
190- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
191- CONVERT_EXCEPTION_TO_ERROR_CODE({ *ret_num_inputs = container->num_inputs(); })
192-}
193- 
194-AOTIRuntimeError AOTInductorModelContainerGetInputName(AOTInductorModelContainerHandle container_handle,
195- size_t input_idx, const char** ret_input_names)
196-{
197- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
198- CONVERT_EXCEPTION_TO_ERROR_CODE({ *ret_input_names = container->input_name(input_idx); })
199-}
200- 
201-AOTIRuntimeError AOTInductorModelContainerGetNumOutputs(AOTInductorModelContainerHandle container_handle,
202- size_t* ret_num_outputs)
203-{
204- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
205- CONVERT_EXCEPTION_TO_ERROR_CODE({ *ret_num_outputs = container->num_outputs(); })
206-}
207- 
208-AOTIRuntimeError AOTInductorModelContainerGetOutputName(AOTInductorModelContainerHandle container_handle,
209- size_t output_idx, const char** ret_output_names)
210-{
211- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
212- CONVERT_EXCEPTION_TO_ERROR_CODE({ *ret_output_names = container->output_name(output_idx); })
213-}
214- 
215-AOTIRuntimeError AOTInductorModelContainerGetCallSpec(AOTInductorModelContainerHandle container_handle,
216- const char** in_spec, const char** out_spec)
217-{
218- auto* container = reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
219- CONVERT_EXCEPTION_TO_ERROR_CODE({
220- *in_spec = container->get_in_spec();
221- *out_spec = container->get_out_spec();
222- })
223-}
224- 
225-AOTIRuntimeError AOTInductorModelCreate(AOTInductorModelHandle* model_handle,
226- AOTInductorConstantMapHandle constant_map_handle)
227-{
228- CONVERT_EXCEPTION_TO_ERROR_CODE({
229- auto constant_map = std::make_shared<torch::aot_inductor::ConstantMap>();
230- auto constant_array = std::make_shared<std::vector<torch::aot_inductor::ConstantHandle> >();
231- auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
232- 
233- auto model = new torch::aot_inductor::AOTInductorModel(
234- constant_map, constant_array,
235- "cpu", // device_str is hardcoded, as AOTInductorModelCreate is only use for CPU models
236- "");
237- 
238- if (input_map) {
239- for (auto const& kv : *input_map) {
240- constant_map->emplace(kv.first, kv.second);
241- }
242- } else {
243- model->load_constants();
244- }
245- 
246- *model_handle = reinterpret_cast<AOTInductorModelHandle>(model);
247- })
248-}
249- 
250-AOTIRuntimeError AOTInductorModelRun(AOTInductorModelHandle model_handle, AtenTensorHandle* input_handles,
251- AtenTensorHandle* output_handles)
252-{
253- auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
254- CONVERT_EXCEPTION_TO_ERROR_CODE({
255- AOTINoGradGuard guard;
256- model->run_impl(input_handles, output_handles, (torch::aot_inductor::DeviceStreamType)nullptr, nullptr);
257- })
258-}
259- 
260-AOTIRuntimeError AOTInductorModelDelete(AOTInductorModelHandle model_handle)
261-{
262- CONVERT_EXCEPTION_TO_ERROR_CODE(
263- {
264- auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
265- delete model;
266- })
267-}
268- 
269-AOTIRuntimeError AOTInductorModelGetNumOutputs(AOTInductorModelHandle model_handle,
270- size_t* ret_num_outputs)
271-{
272- CONVERT_EXCEPTION_TO_ERROR_CODE(
273- {
274- auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
275- *ret_num_outputs = model->num_outputs();
276- })
277-}
278- 
279-AOTIRuntimeError AOTInductorModelUpdateConstantsMap(AOTInductorModelHandle model_handle,
280- AOTInductorConstantMapHandle constant_map_handle)
281-{
282- auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
283- CONVERT_EXCEPTION_TO_ERROR_CODE({
284- auto constant_map = std::make_shared<torch::aot_inductor::ConstantMap>();
285- auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
286- 
287- for (auto const& kv : *input_map) {
288- constant_map->emplace(kv.first, kv.second);
289- }
290- model->update_constants_map(std::move(constant_map));
291- })
292-}
293- 
294-} // extern "C"
Mtorch_npu/_inductor/codegen/cpp_wrapper.py+41-53
@@ -28,19 +28,19 @@ from torch._inductor.codegen.multi_kernel import MultiKernelCall
28from torch._inductor.codegen.wrapper import PythonWrapperCodegen, SymbolicCallArg28from torch._inductor.codegen.wrapper import PythonWrapperCodegen, SymbolicCallArg
29from torch._inductor.ir import GraphPartitionSignature29from torch._inductor.ir import GraphPartitionSignature
30from torch._inductor.runtime.runtime_utils import dynamo_timed30from torch._inductor.runtime.runtime_utils import dynamo_timed
31-from torch._inductor.utils import ALIGN_BYTES, IndentedBuffer31+from torch._inductor.utils import IndentedBuffer
32from torch._inductor.virtualized import V32from torch._inductor.virtualized import V
33 33 
34from .. import config as npu_config34from .. import config as npu_config
35from ..runtime.triton_heuristics import GridExprNpu35from ..runtime.triton_heuristics import GridExprNpu
36-from ..utils import NPU_ALIGN_BYTES, triton_support_ffts, triton_support_auto_blockify36+from ..utils import triton_support_ffts, triton_support_auto_blockify
37- 
38-config.triton.autotune_at_compile_time = False
39- 
40- 
41-# follow triton-ascend implement
42-DTYPE_TO_CPP[torch.bool] = "int32_t"
43 37 
38+# follow triton-ascend implement, except torch.bfloat16
39+# torch.bfloat16 -> "float" will be specially deal in codegen_tensor_item_npu
40+DTYPE_TO_TA_TYPE = {
41+ **DTYPE_TO_CPP,
42+ torch.bool: "int32_t"
43+}
44 44 
45@dataclasses.dataclass45@dataclasses.dataclass
46class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):46class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):
@@ -160,7 +160,7 @@ class DeferredNpuTritonCallWrapper(DeferredTritonCallWrapper):
160 enable_simt = npu_config.is_ascend950 and (160 enable_simt = npu_config.is_ascend950 and (
161 "simt" in params["parallel_mode"] or params["force_simt_only"]161 "simt" in params["parallel_mode"] or params["force_simt_only"]
162 )162 )
163- enable_auto_blockify = not getattr(triton_meta, "has_auto_blockify_blacklist_op", False) and triton_support_auto_blockify()163+ enable_auto_blockify = not params.get("has_auto_blockify_blacklist_op", False) and triton_support_auto_blockify()
164 prefix.splice(f"""164 prefix.splice(f"""
165 auto launch_call = [=]() {{165 auto launch_call = [=]() {{
166 {wrapper.generate_args_decl(prefix, call_args, arg_types, arg_signatures, True, force_simt_only)}166 {wrapper.generate_args_decl(prefix, call_args, arg_types, arg_signatures, True, force_simt_only)}
@@ -219,12 +219,8 @@ class CppWrapperNpu(CppWrapperGpu):
219 219 
220 if V.graph.aot_mode:220 if V.graph.aot_mode:
221 if config.aot_inductor.dynamic_linkage:221 if config.aot_inductor.dynamic_linkage:
222- with open(222+ self.header.splice(self._adapt_community_interface_cpp())
223- os.path.join(223+ self.header.splice("\n")
224- os.path.dirname(__file__), "aoti_runtime", "interface.cpp"
225- )
226- ) as f:
227- self.header.splice(f.read())
228 else:224 else:
229 # we produce a separate model header for each model in static linkage225 # we produce a separate model header for each model in static linkage
230 self.header.splice(f"""#include \"{self.model_class_name_suffix}.h\"""")226 self.header.splice(f"""#include \"{self.model_class_name_suffix}.h\"""")
@@ -246,6 +242,32 @@ class CppWrapperNpu(CppWrapperGpu):
246 maybe_hipify_code_wrapper(self.device_codegen.kernel_driver())242 maybe_hipify_code_wrapper(self.device_codegen.kernel_driver())
247 )243 )
248 244 
245+ @staticmethod
246+ def _adapt_community_interface_cpp() -> str:
247+ """Reuse the upstream interface.cpp from torch and adapt it for NPU.
248+ 
249+ Instead of maintaining a duplicated copy, read the community
250+ interface.cpp and apply the minimal NPU-specific transformations:
251+ - Redirect model_container.h include to the torch_npu variant, since
252+ NPU ships its own model container implementation.
253+ - Default to "npu" instead of "cuda" in AOTInductorModelContainerCreate.
254+ """
255+ community_interface_path = os.path.join(
256+ os.path.dirname(torch.__file__),
257+ "_inductor",
258+ "codegen",
259+ "aoti_runtime",
260+ "interface.cpp",
261+ )
262+ with open(community_interface_path) as f:
263+ content = f.read()
264+ content = content.replace(
265+ "<torch/csrc/inductor/aoti_runtime/model_container.h>",
266+ "<torch_npu/csrc/inductor/aoti_runtime/model_container.h>",
267+ )
268+ content = content.replace("cuda", "npu")
269+ return content
270+ 
249 def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):271 def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
250 expr = f"{kernel_name}_{node.name}_numel"272 expr = f"{kernel_name}_{node.name}_numel"
251 273 
@@ -257,28 +279,6 @@ class CppWrapperNpu(CppWrapperGpu):
257 self.writeline(f"{expr} = {cexpr(numel_expr)};")279 self.writeline(f"{expr} = {cexpr(numel_expr)};")
258 return SymbolicCallArg(expr, numel_expr)280 return SymbolicCallArg(expr, numel_expr)
259 281 
260- def codegen_inputs(self):
261- # See Note: [Input Alignment handling in Inductor]
262- #
263- # JIT Inductor does not guard on input alignment. It relies on copy_misaligned_inputs to
264- # copy misaligned inputs to aligned buffers. For AOTInductor, we expect users to use it
265- # as non-Python deployment for its best performance, so implicitly copying misaligned inputs
266- # to aligned buffers is going to bring a surprising performance hit. Instead, we check input
267- # alignment and throw an error if any input is misaligned.
268- if V.graph.aot_mode and V.graph.inputs_to_check:
269- for idx in V.graph.inputs_to_check:
270- input_name = V.graph.graph_input_names[idx]
271- 
272- self.prefix.splice(
273- f"""
274- if ((long({input_name}.data_ptr()) & ({NPU_ALIGN_BYTES} -1)) != 0) {{
275- throw std::runtime_error("{input_name} is not aligned to {NPU_ALIGN_BYTES} bytes");
276- }}
277- """
278- )
279- 
280- super().codegen_inputs()
281- 
282 def _generate_kernel_call_helper(282 def _generate_kernel_call_helper(
283 self,283 self,
284 kernel_name: str,284 kernel_name: str,
@@ -364,18 +364,6 @@ class CppWrapperNpu(CppWrapperGpu):
364 with dynamo_timed("CppWrapperNpu.generate", log_pt2_compile_event=True):364 with dynamo_timed("CppWrapperNpu.generate", log_pt2_compile_event=True):
365 return super().generate(is_inference)365 return super().generate(is_inference)
366 366 
367- def prepare_triton_kernel_call(self, call_args):
368- new_call_args = call_args
369- if npu_config.inductor_static_mode:
370- # in inductor_static_mode, numel arg is constexpr, remove all Integer constant args from call_args
371- new_call_args = [
372- call_arg
373- for call_arg in call_args
374- if not isinstance(call_arg, sympy.Integer)
375- ]
376- 
377- return super().prepare_triton_kernel_call(new_call_args)
378- 
379 def codegen_tensor_item_npu(367 def codegen_tensor_item_npu(
380 self, dtype: torch.dtype, tensor: str, scalar: str, indented_buffer=None368 self, dtype: torch.dtype, tensor: str, scalar: str, indented_buffer=None
381 ):369 ):
@@ -384,7 +372,7 @@ class CppWrapperNpu(CppWrapperGpu):
384 372 
385 if dtype == torch.float16 or dtype == torch.bfloat16:373 if dtype == torch.float16 or dtype == torch.bfloat16:
386 scalar_tmp = f"{scalar}_tmp"374 scalar_tmp = f"{scalar}_tmp"
387- writer.writeline(f"{DTYPE_TO_CPP[dtype]} {scalar_tmp};")375+ writer.writeline(f"{DTYPE_TO_TA_TYPE[dtype]} {scalar_tmp};")
388 writer.writeline(376 writer.writeline(
389 f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_item_{dtype_str}({tensor}, &{scalar_tmp}));"377 f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_item_{dtype_str}({tensor}, &{scalar_tmp}));"
390 )378 )
@@ -392,12 +380,12 @@ class CppWrapperNpu(CppWrapperGpu):
392 struct_data = f"float {scalar} __attribute__((aligned(4)));"380 struct_data = f"float {scalar} __attribute__((aligned(4)));"
393 arg_data = f"static_cast<float>({scalar})"381 arg_data = f"static_cast<float>({scalar})"
394 else:382 else:
395- writer.writeline(f"{DTYPE_TO_CPP[dtype]} {scalar};")383+ writer.writeline(f"{DTYPE_TO_TA_TYPE[dtype]} {scalar};")
396 writer.writeline(384 writer.writeline(
397 f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_item_{dtype_str}({tensor}, &{scalar}));"385 f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_item_{dtype_str}({tensor}, &{scalar}));"
398 )386 )
399- struct_data = f"{DTYPE_TO_CPP[dtype]} {scalar} __attribute__((aligned(sizeof({DTYPE_TO_CPP[dtype]} ))));"387+ struct_data = f"{DTYPE_TO_TA_TYPE[dtype]} {scalar} __attribute__((aligned(sizeof({DTYPE_TO_TA_TYPE[dtype]} ))));"
400- arg_data = f"static_cast<{DTYPE_TO_CPP[dtype]}>({scalar})"388+ arg_data = f"static_cast<{DTYPE_TO_TA_TYPE[dtype]}>({scalar})"
401 389 
402 return struct_data, arg_data390 return struct_data, arg_data
403 391 
Mtorch_npu/_inductor/codegen/triton.py+6-10
@@ -2071,16 +2071,12 @@ class NPUIndexTritonKernel(TritonKernel):
2071 def gen_numel_args(self, signature, triton_meta, triton_meta_signature, argdefs):2071 def gen_numel_args(self, signature, triton_meta, triton_meta_signature, argdefs):
2072 for node in self.sorted_axis:2072 for node in self.sorted_axis:
2073 arg_name = f"{node.name}_numel"2073 arg_name = f"{node.name}_numel"
2074- if not npu_config.inductor_static_mode:2074+ sizearg = SizeArg(arg_name, node.length)
2075- sizearg = SizeArg(arg_name, node.length)2075+ signature.append(sizearg)
2076- signature.append(sizearg)2076+ triton_meta_signature[arg_name] = signature_of(
2077- triton_meta_signature[arg_name] = signature_of(2077+ sizearg, size_dtype=self.index_dtype
2078- sizearg, size_dtype=self.index_dtype2078+ )
2079- )2079+ argdefs.append(ArgName(arg_name))
2080- argdefs.append(ArgName(arg_name))
2081- else:
2082- argdefs.append(ArgName(arg_name, is_constexpr=True))
2083- triton_meta["constants"][arg_name] = node.length
2084 2080 
2085 # BLOCK and SUB_BLOCK definitions2081 # BLOCK and SUB_BLOCK definitions
2086 def add_autotune_args(self, argdefs, signature, triton_meta_signature):2082 def add_autotune_args(self, argdefs, signature, triton_meta_signature):
Mtorch_npu/_inductor/config.py+1-5
@@ -241,11 +241,7 @@ symbolic_group_allow_templates = tuple(
241 ).split(",")241 ).split(",")
242 if x.strip()242 if x.strip()
243)243)
244-inductor_static_mode = os.environ.get("INDUCTOR_STATIC_MODE", "0").lower() in (244+ 
245- "1",
246- "yes",
247- "true",
248-)
249profile_path = "./profile_result/"245profile_path = "./profile_result/"
250 246 
251fasta_autotune = os.environ.get("FASTAUTOTUNE", "0") == "1"247fasta_autotune = os.environ.get("FASTAUTOTUNE", "0") == "1"
Mtorch_npu/_inductor/utils.py+0-2
@@ -9,8 +9,6 @@ import torch_npu
9import torch._inductor.config as inductor_config9import torch._inductor.config as inductor_config
10log = logging.getLogger("torch._inductor")10log = logging.getLogger("torch._inductor")
11 11 
12-NPU_ALIGN_BYTES = 32
13- 
14 12 
15def get_current_raw_stream(device):13def get_current_raw_stream(device):
16 return torch.npu.current_stream(device).npu_stream14 return torch.npu.current_stream(device).npu_stream