#11004: moreh: use env var for kernel src search path (#12541)

* #11004: Implemented functionality to search for kernels with specified env variable.

* #11004: Added tests

* #11004: Updated function documentation

* #11004: Updated unit tests

* #11004: Fixed TT_FATAL calls

* #11004: Added license

* #11004: Renamed file

* #11004: Added test_kernel_path_env_var to script

tests/scripts/run_cpp_unit_tests.sh

@@ -7,6 +7,11 @@ if [[ -z "$TT_METAL_HOME" ]]; then
     exit 1
 fi
 
+kernel_path="/tmp/kernels"
+mkdir -p $kernel_path
+TT_METAL_KERNEL_PATH=$kernel_path ./build/test/tt_metal/test_kernel_path_env_var
+rm -rf $kernel_path
+
 if [[ ! -z "$TT_METAL_SLOW_DISPATCH_MODE" ]]; then
     ./build/test/tt_metal/unit_tests
     env python tests/scripts/run_tt_metal.py --dispatch-mode slow

tests/tt_metal/tt_metal/CMakeLists.txt

@@ -31,6 +31,7 @@ set (TT_METAL_TESTS_SRCS
     test_core_range_set.cpp
     test_compile_sets_kernel_binaries.cpp
     test_compile_program.cpp
+    test_kernel_path_env_var.cpp
     test_clean_init.cpp
 )
 

tests/tt_metal/tt_metal/test_kernel_path_env_var.cpp

@@ -0,0 +1,134 @@
+// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <gtest/gtest.h>
+
+#include <exception>
+#include <filesystem>
+
+#include "assert.hpp"
+#include "core_coord.h"
+#include "detail/tt_metal.hpp"
+#include "host_api.hpp"
+#include "impl/kernels/data_types.hpp"
+#include "impl/program/program.hpp"
+#include "llrt/rtoptions.hpp"
+#include "tt_cluster_descriptor_types.h"
+
+using namespace tt;
+using namespace tt::tt_metal;
+using namespace tt::llrt;
+
+class CompileProgramWithKernelPathEnvVarFixture : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        this->validate_preconditions();
+
+        const chip_id_t device_id = 0;
+        this->device_ = CreateDevice(device_id);
+        this->program_ = CreateProgram();
+    }
+
+    void TearDown() override { CloseDevice(this->device_); }
+
+    void create_kernel(const string &kernel_file) {
+        CoreCoord core(0, 0);
+        tt_metal::CreateKernel(
+            this->program_,
+            kernel_file,
+            core,
+            tt_metal::DataMovementConfig{.processor = DataMovementProcessor::RISCV_1, .noc = NOC::RISCV_1_default});
+    }
+
+    void setup_kernel_dir(const string &orig_kernel_file, const string &new_kernel_file) {
+        const string &kernel_dir = OptionsG.get_kernel_dir();
+        const std::filesystem::path &kernel_file_path_under_kernel_dir(kernel_dir + new_kernel_file);
+        const std::filesystem::path &dirs_under_kernel_dir = kernel_file_path_under_kernel_dir.parent_path();
+        std::filesystem::create_directories(dirs_under_kernel_dir);
+
+        const string &metal_root = OptionsG.get_root_dir();
+        const std::filesystem::path &kernel_file_path_under_metal_root(metal_root + orig_kernel_file);
+        std::filesystem::copy(kernel_file_path_under_metal_root, kernel_file_path_under_kernel_dir);
+    }
+
+    void cleanup_kernel_dir() {
+        const string &kernel_dir = OptionsG.get_kernel_dir();
+        for (const std::filesystem::directory_entry &entry : std::filesystem::directory_iterator(kernel_dir)) {
+            std::filesystem::remove_all(entry);
+        }
+    }
+
+    Device *device_;
+    Program program_;
+
+   private:
+    void validate_preconditions() {
+        this->validate_env_vars_are_set();
+        this->validate_kernel_dir_is_valid();
+    }
+
+    void validate_env_vars_are_set() {
+        if (!OptionsG.is_root_dir_specified()) {
+            GTEST_SKIP() << "Skipping test: TT_METAL_HOME must be set";
+        }
+        if (!OptionsG.is_kernel_dir_specified()) {
+            GTEST_SKIP() << "Skipping test: TT_METAL_KERNEL_PATH must be set";
+        }
+    }
+
+    void validate_kernel_dir_is_valid() {
+        const string &kernel_dir = llrt::OptionsG.get_kernel_dir();
+        if (!this->does_path_exist(kernel_dir) || !this->is_path_a_directory(kernel_dir) ||
+            !this->is_dir_empty(kernel_dir)) {
+            GTEST_SKIP() << "Skipping test: TT_METAL_KERNEL_PATH must be an existing, empty directory";
+        }
+    }
+
+    bool does_path_exist(const string &path) {
+        const std::filesystem::path &file_path(path);
+        return std::filesystem::exists(file_path);
+    }
+
+    bool is_path_a_directory(const string &path) {
+        TT_FATAL(this->does_path_exist(path), "{} does not exist", path);
+        const std::filesystem::path &file_path(path);
+        return std::filesystem::is_directory(file_path);
+    }
+
+    bool is_dir_empty(const string &path) {
+        TT_FATAL(this->does_path_exist(path), "{} does not exist", path);
+        TT_FATAL(this->is_path_a_directory(path), "{} is not a directory", path);
+        const std::filesystem::path &file_path(path);
+        return std::filesystem::is_empty(file_path);
+    }
+};
+
+TEST_F(CompileProgramWithKernelPathEnvVarFixture, KernelUnderMetalRootDir) {
+    const string &kernel_file = "tests/tt_metal/tt_metal/test_kernels/dataflow/reader_unary_push_4.cpp";
+    create_kernel(kernel_file);
+    detail::CompileProgram(this->device_, this->program_);
+}
+
+TEST_F(CompileProgramWithKernelPathEnvVarFixture, KernelUnderKernelRootDir) {
+    const string &orig_kernel_file = "tests/tt_metal/tt_metal/test_kernels/dataflow/reader_unary_push_4.cpp";
+    const string &new_kernel_file = "tests/tt_metal/tt_metal/test_kernels/dataflow/new_kernel.cpp";
+    this->setup_kernel_dir(orig_kernel_file, new_kernel_file);
+    this->create_kernel(new_kernel_file);
+    detail::CompileProgram(this->device_, this->program_);
+    this->cleanup_kernel_dir();
+}
+
+TEST_F(CompileProgramWithKernelPathEnvVarFixture, KernelUnderMetalRootDirAndKernelRootDir) {
+    const string &kernel_file = "tests/tt_metal/tt_metal/test_kernels/dataflow/reader_unary_push_4.cpp";
+    this->setup_kernel_dir(kernel_file, kernel_file);
+    this->create_kernel(kernel_file);
+    detail::CompileProgram(this->device_, this->program_);
+    this->cleanup_kernel_dir();
+}
+
+TEST_F(CompileProgramWithKernelPathEnvVarFixture, NonExistentKernel) {
+    const string &kernel_file = "tests/tt_metal/tt_metal/test_kernels/dataflow/non_existent_kernel.cpp";
+    this->create_kernel(kernel_file);
+    EXPECT_THROW(detail::CompileProgram(this->device_, this->program_), std::exception);
+}

tt_metal/host_api.hpp

@@ -114,12 +114,12 @@ Program CreateProgram();
  *
  * Return value: Kernel ID (uintptr_t)
  *
- * | Argument     | Description                                                                                                                          | Type                                                     | Valid Range | Required |
- * |--------------|--------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------|-------------|----------|
- * | program      | The program to which this kernel will be added to                                                                                    | Program &                                                |             | Yes      |
- * | file_name    | Path to kernel src. Assumed to be absolute/relative to CWD, but will fall back to relative path from TT_METAL_HOME.                  | const std::string &                                      |             | Yes      |
- * | core_spec    | Either a single logical core, a range of logical cores or a set of logical core ranges that indicate which cores kernel is placed on | const std::variant<CoreCoord, CoreRange, CoreRangeSet> & |             | Yes      |
- * | config       | Config for data movement or compute kernel                                                                                           | const std::variant<DataMovementConfig,ComputeConfig,EthernetConfig> &   |             | No       |
+ * | Argument     | Description                                                                                                                                 | Type                                                     | Valid Range | Required |
+ * |--------------|---------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------|-------------|----------|
+ * | program      | The program to which this kernel will be added to                                                                                           | Program &                                                |             | Yes      |
+ * | file_name    | Path to kernel src. Assumed to be absolute/relative to CWD, but will fall back to relative path from TT_METAL_HOME/TT_METAL_KERNEL_PATH.    | const std::string &                                      |             | Yes      |
+ * | core_spec    | Either a single logical core, a range of logical cores or a set of logical core ranges that indicate which cores kernel is placed on        | const std::variant<CoreCoord, CoreRange, CoreRangeSet> & |             | Yes      |
+ * | config       | Config for data movement or compute kernel                                                                                                  | const std::variant<DataMovementConfig,ComputeConfig,EthernetConfig> &   |             | No       |
  */
 KernelHandle CreateKernel(
     Program &program,