[v2.10.0][feature]support environment_variable TORCH_NPU_CPU

Atest/npu/test_cpu_fallback_control.py+211-0

@@ -0,0 +1,211 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd
	2	+# Licensed under the BSD 3-Clause License.
	3	+
	4	+import os
	5	+import subprocess
	6	+import sys
	7	+import unittest
	8	+
	9	+import torch
	10	+
	11	+from torch_npu.testing.testcase import TestCase, run_tests
	12	+
	13	+
	14	+CPU_FALLBACK_ENV = "TORCH_NPU_FALLBACK_CPU_DISABLE"
	15	+SUBPROCESS_TIMEOUT = 120
	16	+
	17	+
	18	+# OptionsManager caches the environment variable in a function-local static.
	19	+# Every case must therefore run in a fresh process, with the environment set
	20	+# before torch_npu is imported.
	21	+_CHILD_SCRIPT = r"""
	22	+import sys
	23	+
	24	+import torch
	25	+
	26	+
	27	+def assert_npu_tensor(tensor, name):
	28	+ if tensor.device.type != "npu":
	29	+ raise AssertionError(f"{name} must be on NPU, got {tensor.device}")
	30	+
	31	+
	32	+def run_dispatcher_fmax_out():
	33	+ x = torch.tensor([1.0, float("nan"), 3.0], dtype=torch.float32, device="npu")
	34	+ y = torch.tensor([2.0, 4.0, float("nan")], dtype=torch.float32, device="npu")
	35	+ expected = torch.tensor([2.0, 4.0, 3.0], dtype=torch.float32)
	36	+ for _ in range(2):
	37	+ out = torch.empty_like(x)
	38	+ torch.fmax(x, y, out=out)
	39	+ assert_npu_tensor(out, "fmax.out result")
	40	+ torch.testing.assert_close(out.cpu(), expected)
	41	+
	42	+
	43	+def make_sparse_csr():
	44	+ crow = torch.tensor([0, 2, 4], dtype=torch.int64, device="npu")
	45	+ col = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="npu")
	46	+ values = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device="npu")
	47	+ return torch.sparse_csr_tensor(crow, col, values, size=(2, 2))
	48	+
	49	+
	50	+def run_sparse_csr(reduction):
	51	+ value = make_sparse_csr()
	52	+ op = getattr(torch.ops.aten, f"_sparse_csr_{reduction}").dim_dtype
	53	+ for _ in range(2):
	54	+ result = op(value, [1], True, dtype=None)
	55	+ assert_npu_tensor(result, f"sparse_csr_{reduction} result")
	56	+ # Materialize the result to ensure the fallback and NPU copy-back have
	57	+ # completed. No CPU reference reduction is run before the target op.
	58	+ result.cpu()
	59	+
	60	+
	61	+def run_normal_npu():
	62	+ x = torch.arange(8, dtype=torch.float32, device="npu")
	63	+ y = torch.ones(8, dtype=torch.float32, device="npu")
	64	+ expected = torch.arange(8, dtype=torch.float32) + 1
	65	+ for _ in range(2):
	66	+ result = torch.add(x, y)
	67	+ assert_npu_tensor(result, "add result")
	68	+ torch.testing.assert_close(result.cpu(), expected)
	69	+
	70	+
	71	+torch.npu.set_device(0)
	72	+case = sys.argv[1]
	73	+if case == "dispatcher_fmax_out":
	74	+ run_dispatcher_fmax_out()
	75	+elif case == "sparse_csr_sum":
	76	+ run_sparse_csr("sum")
	77	+elif case == "sparse_csr_prod":
	78	+ run_sparse_csr("prod")
	79	+elif case == "normal_npu":
	80	+ run_normal_npu()
	81	+else:
	82	+ raise AssertionError(f"unknown case: {case}")
	83	+
	84	+torch.npu.synchronize()
	85	+print("CASE_SUCCESS")
	86	+"""
	87	+
	88	+
	89	+@unittest.skipIf(not torch.npu.is_available(), "requires NPU")
	90	+class TestCpuFallbackControl(TestCase):
	91	+ @staticmethod
	92	+ def _run_case(case, env_value):
	93	+ env = os.environ.copy()
	94	+ if env_value is None:
	95	+ env.pop(CPU_FALLBACK_ENV, None)
	96	+ else:
	97	+ env[CPU_FALLBACK_ENV] = env_value
	98	+ return subprocess.run(
	99	+ [sys.executable, "-c", _CHILD_SCRIPT, case],
	100	+ env=env,
	101	+ capture_output=True,
	102	+ text=True,
	103	+ timeout=SUBPROCESS_TIMEOUT,
	104	+ check=False,
	105	+ )
	106	+
	107	+ @staticmethod
	108	+ def _output(result):
	109	+ return (result.stdout or "") + (result.stderr or "")
	110	+
	111	+ def _assert_fallback_allowed(self, case, env_value, warning_needle, op_name=None):
	112	+ result = self._run_case(case, env_value)
	113	+ output = self._output(result)
	114	+ self.assertEqual(
	115	+ result.returncode,
	116	+ 0,
	117	+ f"fallback should be allowed for {case}, env={env_value!r}\n{output}",
	118	+ )
	119	+ self.assertIn("CASE_SUCCESS", output)
	120	+ self.assertEqual(
	121	+ output.lower().count(warning_needle.lower()),
	122	+ 1,
	123	+ f"fallback warning must be emitted once for {case}\n{output}",
	124	+ )
	125	+ if op_name is not None:
	126	+ self.assertIn(op_name, output)
	127	+
	128	+ def _assert_fallback_blocked(self, case, warning_needle, op_name=None):
	129	+ result = self._run_case(case, "1")
	130	+ output = self._output(result)
	131	+ self.assertNotEqual(
	132	+ result.returncode,
	133	+ 0,
	134	+ f"fallback should be blocked for {case}\n{output}",
	135	+ )
	136	+ self.assertNotIn("CASE_SUCCESS", output)
	137	+ self.assertIn(CPU_FALLBACK_ENV, output)
	138	+ self.assertNotIn(
	139	+ warning_needle.lower(),
	140	+ output.lower(),
	141	+ f"strict mode must fail before the fallback warning for {case}\n{output}",
	142	+ )
	143	+ if op_name is not None:
	144	+ self.assertIn(op_name, output)
	145	+
	146	+ def test_dispatcher_fallback_default_and_explicitly_allowed(self):
	147	+ for env_value in (None, "0"):
	148	+ with self.subTest(env_value=env_value):
	149	+ self._assert_fallback_allowed(
	150	+ "dispatcher_fmax_out",
	151	+ env_value,
	152	+ "will fall back to run on the CPU",
	153	+ "aten::fmax.out",
	154	+ )
	155	+
	156	+ def test_dispatcher_fallback_disabled(self):
	157	+ self._assert_fallback_blocked(
	158	+ "dispatcher_fmax_out",
	159	+ "will fall back to run on the CPU",
	160	+ "aten::fmax.out",
	161	+ )
	162	+
	163	+ def test_sparse_csr_sum_fallback_default_and_explicitly_allowed(self):
	164	+ for env_value in (None, "0"):
	165	+ with self.subTest(env_value=env_value):
	166	+ self._assert_fallback_allowed(
	167	+ "sparse_csr_sum",
	168	+ env_value,
	169	+ "will fall back to CPU",
	170	+ "aten::_sparse_csr_sum.dim_dtype",
	171	+ )
	172	+
	173	+ def test_sparse_csr_sum_fallback_disabled(self):
	174	+ self._assert_fallback_blocked(
	175	+ "sparse_csr_sum",
	176	+ "will fall back to CPU",
	177	+ "aten::_sparse_csr_sum.dim_dtype",
	178	+ )
	179	+
	180	+ def test_sparse_csr_prod_fallback_default_and_explicitly_allowed(self):
	181	+ for env_value in (None, "0"):
	182	+ with self.subTest(env_value=env_value):
	183	+ self._assert_fallback_allowed(
	184	+ "sparse_csr_prod",
	185	+ env_value,
	186	+ "will fall back to CPU",
	187	+ "aten::_sparse_csr_prod.dim_dtype",
	188	+ )
	189	+
	190	+ def test_sparse_csr_prod_fallback_disabled(self):
	191	+ self._assert_fallback_blocked(
	192	+ "sparse_csr_prod",
	193	+ "will fall back to CPU",
	194	+ "aten::_sparse_csr_prod.dim_dtype",
	195	+ )
	196	+
	197	+ def test_normal_npu_kernel_is_not_blocked(self):
	198	+ result = self._run_case("normal_npu", "1")
	199	+ output = self._output(result)
	200	+ self.assertEqual(
	201	+ result.returncode,
	202	+ 0,
	203	+ f"strict mode must not block a normal NPU kernel\n{output}",
	204	+ )
	205	+ self.assertIn("CASE_SUCCESS", output)
	206	+ self.assertNotIn("will fall back", output.lower())
	207	+ self.assertNotIn("fallback to run on the cpu", output.lower())
	208	+
	209	+
	210	+if __name__ == "__main__":
	211	+ run_tests()

Mtorch_npu/csrc/aten/VariableFallbackKernel.cpp+8-1

	@@ -9,6 +9,7 @@
9	#include <ATen/core/ivalue.h>	9	#include <ATen/core/ivalue.h>
10		10
11	#include "torch_npu/csrc/core/npu/NPUException.h"	11	#include "torch_npu/csrc/core/npu/NPUException.h"
		12	+#include "torch_npu/csrc/framework/utils/CpuFallbackUtils.h"
12		13
13	/*	14	/*
14	* This file implements a variable fallback kernel for custom operators.	15	* This file implements a variable fallback kernel for custom operators.
	@@ -246,7 +247,13 @@ bool has_op_name_warned(const std::string& op_name)
246		247
247	void npu_cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack)	248	void npu_cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack)
248	{	249	{
249	- if (!has_op_name_warned(c10::toString(op.schema().operator_name()))) {	250	+ const auto op_name = c10::toString(op.schema().operator_name());
		251	+ at_npu::native::CheckCpuFallbackAllowed(
		252	+ op_name,
		253	+ at_npu::native::CpuFallbackKind::Dispatcher,
		254	+ "no kernel is registered for the NPU PrivateUse1 backend");
		255	+
		256	+ if (!has_op_name_warned(op_name)) {
250	TORCH_NPU_WARN("CAUTION: The operator '",	257	TORCH_NPU_WARN("CAUTION: The operator '",
251	op.schema().operator_name(),	258	op.schema().operator_name(),
252	"' is not currently supported ",	259	"' is not currently supported ",

Mtorch_npu/csrc/aten/ops/SparseCsrReductionFallbackKernelNpu.cpp+21-0

	@@ -3,6 +3,7 @@
3	#include <ATen/ops/_sparse_csr_prod.h>	3	#include <ATen/ops/_sparse_csr_prod.h>
4	#include <ATen/ops/_sparse_csr_sum.h>	4	#include <ATen/ops/_sparse_csr_sum.h>
5	#include <torch/library.h>	5	#include <torch/library.h>
		6	+#include "torch_npu/csrc/framework/utils/CpuFallbackUtils.h"
6		7
7	#include <optional>	8	#include <optional>
8		9
	@@ -13,6 +14,16 @@ at::Tensor sparse_csr_sum_cpu_fallback(
13	at::IntArrayRef dim,	14	at::IntArrayRef dim,
14	bool keepdim,	15	bool keepdim,
15	std::optional<at::ScalarType> dtype) {	16	std::optional<at::ScalarType> dtype) {
		17	+ at_npu::native::CheckCpuFallbackAllowed(
		18	+ "aten::_sparse_csr_sum.dim_dtype",
		19	+ at_npu::native::CpuFallbackKind::ExplicitKernel,
		20	+ "torch_npu currently provides only a CPU implementation for this "
		21	+ "Sparse CSR reduction");
		22	+ TORCH_NPU_WARN_ONCE(
		23	+ "CPU fallback is allowed because "
		24	+ "TORCH_NPU_FALLBACK_CPU_DISABLE=0. Since the NPU interface does not "
		25	+ "support operator 'aten::_sparse_csr_sum.dim_dtype', the operator will "
		26	+ "fall back to CPU.");
16	const at::OptionalDeviceGuard device_guard(at::device_of(self));	27	const at::OptionalDeviceGuard device_guard(at::device_of(self));
17	at::Tensor cpu_self = self.cpu();	28	at::Tensor cpu_self = self.cpu();
18	at::Tensor cpu_result = at::_sparse_csr_sum(cpu_self, dim, keepdim, dtype);	29	at::Tensor cpu_result = at::_sparse_csr_sum(cpu_self, dim, keepdim, dtype);
	@@ -24,6 +35,16 @@ at::Tensor sparse_csr_prod_cpu_fallback(
24	at::IntArrayRef dim,	35	at::IntArrayRef dim,
25	bool keepdim,	36	bool keepdim,
26	std::optional<at::ScalarType> dtype) {	37	std::optional<at::ScalarType> dtype) {
		38	+ at_npu::native::CheckCpuFallbackAllowed(
		39	+ "aten::_sparse_csr_prod.dim_dtype",
		40	+ at_npu::native::CpuFallbackKind::ExplicitKernel,
		41	+ "torch_npu currently provides only a CPU implementation for this "
		42	+ "Sparse CSR reduction");
		43	+ TORCH_NPU_WARN_ONCE(
		44	+ "CPU fallback is allowed because "
		45	+ "TORCH_NPU_FALLBACK_CPU_DISABLE=0. Since the NPU interface does not "
		46	+ "support operator 'aten::_sparse_csr_prod.dim_dtype', the operator will "
		47	+ "fall back to CPU.");
27	const at::OptionalDeviceGuard device_guard(at::device_of(self));	48	const at::OptionalDeviceGuard device_guard(at::device_of(self));
28	at::Tensor cpu_self = self.cpu();	49	at::Tensor cpu_self = self.cpu();
29	at::Tensor cpu_result = at::_sparse_csr_prod(cpu_self, dim, keepdim, dtype);	50	at::Tensor cpu_result = at::_sparse_csr_prod(cpu_self, dim, keepdim, dtype);

Mtorch_npu/csrc/core/npu/register/OptionsManager.cpp+9-0

	@@ -61,6 +61,15 @@ bool OptionsManager::IsResumeModeEnable()
61	return isResumeModeEnable;	61	return isResumeModeEnable;
62	}	62	}
63		63
		64	+bool OptionsManager::IsCpuFallbackEnable()
		65	+{
		66	+ const static bool isCpuFallbackEnable = []() -> bool {
		67	+ int32_t disable = OptionsManager::GetBoolTypeOption("TORCH_NPU_FALLBACK_CPU_DISABLE", 0);
		68	+ return disable == 0;
		69	+ }();
		70	+ return isCpuFallbackEnable;
		71	+}
		72	+
64	bool OptionsManager::IsSubCommRootInfoEnable()	73	bool OptionsManager::IsSubCommRootInfoEnable()
65	{	74	{
66	const static bool isSubCommRootInfoEnable = []() -> bool {	75	const static bool isSubCommRootInfoEnable = []() -> bool {

Mtorch_npu/csrc/core/npu/register/OptionsManager.h+1-0

	@@ -106,6 +106,7 @@ class OptionsManager {
106	public:	106	public:
107	static bool IsHcclZeroCopyEnable();	107	static bool IsHcclZeroCopyEnable();
108	static bool IsResumeModeEnable();	108	static bool IsResumeModeEnable();
		109	+ static bool IsCpuFallbackEnable();
109	static bool IsSubCommRootInfoEnable();	110	static bool IsSubCommRootInfoEnable();
110	static ReuseMode GetMultiStreamMemoryReuse();	111	static ReuseMode GetMultiStreamMemoryReuse();
111	static bool CheckInfNanModeEnable();	112	static bool CheckInfNanModeEnable();

Atorch_npu/csrc/framework/utils/CpuFallbackUtils.h+38-0

@@ -0,0 +1,38 @@
	1	+#pragma once
	2	+
	3	+#include <string>
	4	+
	5	+#include "torch_npu/csrc/core/npu/NPUException.h"
	6	+#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
	7	+
	8	+namespace at_npu {
	9	+namespace native {
	10	+
	11	+enum class CpuFallbackKind {
	12	+ Dispatcher,
	13	+ ExplicitKernel,
	14	+};
	15	+
	16	+inline void CheckCpuFallbackAllowed(
	17	+ const std::string& op_name,
	18	+ CpuFallbackKind kind,
	19	+ const std::string& reason = "") {
	20	+ if (c10_npu::option::OptionsManager::IsCpuFallbackEnable()) {
	21	+ return;
	22	+ }
	23	+
	24	+ const char* fallback_kind = kind == CpuFallbackKind::Dispatcher ? "dispatcher backend fallback" : "explicit CPU kernel";
	25	+ TORCH_CHECK(
	26	+ false,
	27	+ "The operator '",
	28	+ op_name,
	29	+ "' would execute its main computation on CPU through ",
	30	+ fallback_kind,
	31	+ reason.empty() ? "" : " because ",
	32	+ reason,
	33	+ ", but CPU fallback is disabled by TORCH_NPU_FALLBACK_CPU_DISABLE=1.",
	34	+ OPS_ERROR(ErrCode::NOT_SUPPORT));
	35	+}
	36	+
	37	+} // namespace native
	38	+} // namespace at_npu