From bf39209bd16baf5b92bbf74eaa32fca513ecd07d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E6=A1=82=E5=86=9B?= <huangguijun@huawei.com>
Date: Mon, 17 Mar 2025 15:21:00 +0000
Subject: [PATCH] =?UTF-8?q?!2943=20check=20varialbe=5Fuse=5F1g=5Fhuge=5Fpa?=
 =?UTF-8?q?ge=20value=20Merge=20pull=20request=20!2943=20from=20=E9=BB=84?=
 =?UTF-8?q?=E6=A1=82=E5=86=9B/0317=5Foption=5Fcheck?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../tests/st/util/testcase/npu_attrs_test.cc     | 16 ++++++++++++++++
 .../tests/ut/util/testcase/npu_attrs_test.cc     | 16 ++++++++++++++++
 tf_adapter/util/npu_attrs.cc                     |  2 ++
 3 files changed, 34 insertions(+)

diff --git a/tf_adapter/tests/st/util/testcase/npu_attrs_test.cc b/tf_adapter/tests/st/util/testcase/npu_attrs_test.cc
index 408580c90..032af0610 100644
--- a/tf_adapter/tests/st/util/testcase/npu_attrs_test.cc
+++ b/tf_adapter/tests/st/util/testcase/npu_attrs_test.cc
@@ -150,6 +150,22 @@ TEST_F(NpuAttrTest, CheckPrecisionModeV2) {
   EXPECT_EQ(s.ok(), false);
 }
 
+TEST_F(NpuAttrTest, CheckVariableUse1gHugePageInvalid) {
+  GraphOptimizationPassOptions options;
+  SessionOptions session_options;
+  session_options.config.mutable_graph_options()->mutable_optimizer_options()->set_do_function_inlining(true);
+  auto *custom_config =
+      session_options.config.mutable_graph_options()->mutable_rewrite_options()->add_custom_optimizers();
+  custom_config->set_name("NpuOptimizer");
+  options.session_options = &session_options;
+
+  AttrValue precision_mode_v2 = AttrValue();
+  precision_mode_v2.set_i(3);
+  (*custom_config->mutable_parameter_map())["variable_use_1g_huge_page"] = precision_mode_v2;
+  Status s = NpuAttrs::SetNpuOptimizerAttr(options, reinterpret_cast<Node *>(1));
+  EXPECT_EQ(s.ok(), false);
+}
+
 TEST_F(NpuAttrTest, CheckModifyMixList_Failed) {
   GraphOptimizationPassOptions options;
   SessionOptions session_options;
diff --git a/tf_adapter/tests/ut/util/testcase/npu_attrs_test.cc b/tf_adapter/tests/ut/util/testcase/npu_attrs_test.cc
index 58494d316..0f4001651 100644
--- a/tf_adapter/tests/ut/util/testcase/npu_attrs_test.cc
+++ b/tf_adapter/tests/ut/util/testcase/npu_attrs_test.cc
@@ -241,6 +241,22 @@ TEST_F(NpuAttrTest, CheckPrecisionModeV2) {
   EXPECT_EQ(s.ok(), false);
 }
 
+TEST_F(NpuAttrTest, CheckVariableUse1gHugePageInvalid) {
+  GraphOptimizationPassOptions options;
+  SessionOptions session_options;
+  session_options.config.mutable_graph_options()->mutable_optimizer_options()->set_do_function_inlining(true);
+  auto *custom_config =
+      session_options.config.mutable_graph_options()->mutable_rewrite_options()->add_custom_optimizers();
+  custom_config->set_name("NpuOptimizer");
+  options.session_options = &session_options;
+
+  AttrValue precision_mode_v2 = AttrValue();
+  precision_mode_v2.set_i(3);
+  (*custom_config->mutable_parameter_map())["variable_use_1g_huge_page"] = precision_mode_v2;
+  Status s = NpuAttrs::SetNpuOptimizerAttr(options, reinterpret_cast<Node *>(1));
+  EXPECT_EQ(s.ok(), false);
+}
+
 TEST_F(NpuAttrTest, CheckPrecisionModeV2_Failed_WhenAssignedBoth) {
   GraphOptimizationPassOptions options;
   SessionOptions session_options;
diff --git a/tf_adapter/util/npu_attrs.cc b/tf_adapter/util/npu_attrs.cc
index 8e580ab6e..f1e23afb0 100644
--- a/tf_adapter/util/npu_attrs.cc
+++ b/tf_adapter/util/npu_attrs.cc
@@ -2383,6 +2383,8 @@ Status NpuAttrs::SetNpuOptimizerAttr(const GraphOptimizationPassOptions &options
       }
       if (params.count("variable_use_1g_huge_page") > 0) {
         variable_use_1g_huge_page = std::to_string(params.at("variable_use_1g_huge_page").i());
+        const static std::vector<int32_t> kAllowList = {0, 1, 2};
+        NPU_REQUIRES_OK(CheckValueAllowed<int32_t>(params.at("variable_use_1g_huge_page").i(), kAllowList));
       }
       if (params.count("memory_optimization_policy") > 0) {
         memory_optimization_policy = params.at("memory_optimization_policy").s();
-- 
Gitee