Env based cpu isa/feat mask

src/cpu.cpp

@@ -23,6 +23,7 @@
 #endif // __wasi__
 #include <stdio.h>
 #include <string.h>
+#include <iostream>
 
 #ifdef _OPENMP
 #if NCNN_SIMPLEOMP
@@ -617,8 +618,14 @@ static int get_cpu_support_x86_avx_vnni()
     return cpu_info[0] & (1u << 4);
 }
 
+static int g_force_avx512_disabled = 0;
 static int get_cpu_support_x86_avx512()
 {
+    if (g_force_avx512_disabled == 1)
+    {
+        std::cerr << "AVX512 support is disabled due to environment variable setting." << std::endl;
+        return 0;
+    }
 #if __APPLE__
     return get_hw_capability("hw.optional.avx512f")
            && get_hw_capability("hw.optional.avx512bw")
@@ -1868,6 +1875,14 @@ static void initialize_global_cpu_info()
     g_powersave = 0;
     initialize_cpu_thread_affinity_mask(g_cpu_affinity_mask_all, g_cpu_affinity_mask_little, g_cpu_affinity_mask_big);
 
+    // ��黷������������ g_force_avx512_disabled
+    const char* env_ncnn_x86_avx512 = std::getenv("NCNN_X86_AVX512");
+    if (env_ncnn_x86_avx512 && atoi(env_ncnn_x86_avx512) == 0)
+    {
+        g_force_avx512_disabled = 1;
+        std::cerr << "AVX512 support is disabled due to environment variable setting." << std::endl;
+    }
+
 #if (defined _WIN32 && (__aarch64__ || __arm__))
     if (!is_being_debugged())
     {
@@ -2145,6 +2160,12 @@ int cpu_support_arm_vfpv4()
 int cpu_support_arm_asimdhp()
 {
     try_initialize_global_cpu_info();
+    // ��黷������ NCNN_ISA
+    const char* ncnnIsaEnv = std::getenv("NCNN_ISA");
+    if (ncnnIsaEnv && strstr(ncnnIsaEnv, "+asimdhp") != nullptr)
+    {
+        return true;
+    }
 #if __aarch64__
 #if defined _WIN32
     return g_cpu_support_arm_asimdhp;

tools/pnnx/src/save_ncnn.cpp

@@ -462,7 +462,85 @@ int save_ncnn(const Graph& g, const std::string& parampath, const std::string& b
 
     fclose(pyfp);
 
+    // Generate C++ test
+    std::string cpppath = pypath.substr(0, pypath.find_last_of('.')) + ".cpp";
+    FILE* cppfp = fopen(cpppath.c_str(), "wb");
+    if (!cppfp)
+    {
+        fprint(stderr, "fopen %s failed\n", cpppath.c_str());
+        return -1;
+    }
+
+    fprintf(cppfp, "#include <iostream>\n");
+    fprintf(cppfp, "#include <vector>\n");
+    fprintf(cppfp, "#include <random>\n");
+    fprintf(cppfp, "#include \"ncnn.hpp\"\n\n");
+
+    fprintf(cppfp, "void test_inference(const std::string& parampath, const std::string& binpath)\n");
+    fprintf(cppfp, "{\n");
+    fprintf(cppfp, "    ncnn:::Net net;\n");
+    fprintf(cppfp, "    net.load_param(parampath.c_str());\n");
+    fprintf(cppfp, "    ner.load_model(binpath.c_str());\n\n");
+
+    for (int input_index = 0;; input_index++)
+    {
+        std::string input_name = std::string("in") + std::to_string(input_index);
+        const Operand* r = g.get_operand(input_name);
+        if (!r)
+            break;
+
+        fprintf(cppfp, "    std::vector<float> %s(", input_name.c_str());
+        for (size_t i = 0; i < r->shape.size(); i++) 
+        {
+            fprintf(cppfp, "%d", r->shape[i]);
+            if (i + 1 != r->shape.size())
+                fprintf(cppfp, " * ");
+        }
+        fprint(cppfp, ");\n");
+
+        if (type_is_integer(r->type))
+        {
+            fprintf(cppfp, "    std::default_random_engine engine;\n");
+            fprintf(cppfp, "    std::uniform_int_distribution<int> dist(0, 9);\n");
+            fprintf(cppfp, "    for (float& v : %s) {{ v = dist(engine); }}\n", input_name.c_str());
+
+        }
+        else
+        {
+            fprintf(cppfp, "    std::default_random_engine engine;\n");
+            fprintf(cppfp, "    std::uniform_real_distribution<float> dist(0.0, 1.0);\n");
+            fprintf(cppfp, "    for (float& v : %s) {{ v = dist(engine); }}\n", input_name.c_str());
+        }
+
+        fprintf(cppfp, "    ncnn::Mat %s_mat = ncnn::Mat::from_pixels(%s.data(), ncnn::Mat::PixelType::PIXEL_GRAY, %d, %d);\n",
+                input_name.c_str(), input_name.c_str(), r->shape[2], r->shape[1]);
+        // Assuming shape order: batch, height, width
+    fprintf(cppfp, "    net.input(\"net_input\", %s_mat);\n", input_name.c_str());
+
+    }
+    fprintf(cppfp, "    std::vector<ncnn::Mat> output;\n");
+    fprintf(cppfp, "    net.extract(\"output\", output);\n\n");
+
+    fprintf(cppfp, "    for (const auto& m : output) {\n");
+    fprintf(cppfp, "        const float* data = m.channel_data(0);\n");
+    fprintf(cppfp, "        for (int i = 0; i < m.w; i++) {\n");
+    fprintf(cppfp, "            std::cout << data[i] << \" \";\n");
+    fprintf(cppfp, "        }\n");
+    fprintf(cppfp, "        std::cout << std::endl;\n");
+    fprintf(cppfp, "    }\n");
+
+    fprintf(cppfp, "}\n\n");
+
+    fprintf(cppfp, "int main()\n");
+    fprintf(cppfp, "{\n");
+    fprintf(cppfp, "    std::string parampath = \"%s\";\n", parampath.c_str());
+    fprintf(cppfp, "    std::string binpath = \"%s\";\n", binpath.c_str());
+    fprintf(cppfp, "    test_inference(parampath, binpath);\n");
+    fprintf(cppfp, "    return 0;\n");
+    fprintf(cppfp, "}\n");
+
+    fclose(cppfp);
+
     return 0;
-}
 
 } // namespace pnnx

tools/pnnx/tests/ncnn/CMakeLists.txt

@@ -3,6 +3,7 @@ find_package(Python3 REQUIRED COMPONENTS Interpreter)
 
 macro(pnnx_ncnn_add_test name)
     add_test(NAME test_ncnn_${name} COMMAND ${CMAKE_COMMAND} -DPYTHON_EXECUTABLE=${Python3_EXECUTABLE} -DPYTHON_SCRIPT=${CMAKE_CURRENT_SOURCE_DIR}/test_${name}.py -P ${CMAKE_CURRENT_SOURCE_DIR}/../run_test.cmake)
+    add_test(NAME test_cpp_${name} COMMAND ${CMAKE_COMMAND} -DCPP_EXECUTABLE=${Cpp_EXECUTABLE} -DCPP_SCRIPT=${CMAKE_CURRENT_SOURCE_DIR}/test_${name}.cpp -P ${CMAKE_CURRENT_SOURCE_DIR}/../run_test.cmake)
 endmacro()
 
 pnnx_ncnn_add_test(F_adaptive_avg_pool1d)