fix module hierarchy traverse && impl torch-like MACROs

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore/nn/module.hpp

@@ -3,10 +3,13 @@
 #include "parameter.hpp"
 
 #include <unordered_map>
+#include <type_traits>
 
 namespace infinicore::nn {
 class Module {
 public:
+    Module() : path_("") {}
+
     const std::unordered_map<std::string, Parameter> &state_dict() const;
 
     void load_state_dict(const std::unordered_map<std::string, Tensor> &_state_dict);
@@ -17,37 +20,106 @@ class Module {
 
     Tensor register_parameter(const std::string &name, Parameter param);
 
-    // Create a Linear submodule-like parameter set (weight and optional bias)
-    // Mirrors torch.nn.Linear shapes: weight [out_features, in_features], bias [out_features]
-    void linear(const std::string &name, size_t in_features, size_t out_features, bool bias = true);
+    // Get the full path of this module in the hierarchy
+    const std::string &path() const { return path_; }
 
+    // Add an existing submodule to this module's hierarchy
+    // Template parameter M must be a type derived from Module
+    // Returns the submodule for convenience (allows method chaining)
     template <typename M>
     std::shared_ptr<M> add_module(const std::string &name, std::shared_ptr<M> submodule) {
+        // Ensure M is derived from Module (compile-time check)
+        static_assert(std::is_base_of<Module, M>::value,
+                      "Template parameter M must be derived from infinicore::nn::Module");
+
+        // Set the submodule's path based on this module's path
+        std::string submodule_path = path_.empty() ? name : path_ + "." + name;
+        submodule->path_ = submodule_path;
+
+        // Recursively update paths of all descendants
+        // This handles the case where submodules are added in any order
+        submodule->update_submodule_paths();
+
+        // Store in the submodules map (std::shared_ptr<M> automatically converts to std::shared_ptr<Module>)
         submodules_[name] = submodule;
+
         return submodule;
     }
 
+    // Create and register a new submodule by constructing it with the given arguments
+    // Template parameter M must be a type derived from Module
+    // Args are forwarded to M's constructor
     template <typename M, typename... Args>
     std::shared_ptr<M> register_module(const std::string &name, Args &&...args) {
+        // Ensure M is derived from Module (compile-time check)
+        static_assert(std::is_base_of<Module, M>::value,
+                      "Template parameter M must be derived from infinicore::nn::Module");
+
+        // Construct the submodule
         auto submodule = std::make_shared<M>(std::forward<Args>(args)...);
+
         return add_module(name, submodule);
     }
 
+    // Create and register multiple submodules of the same type
+    // Each submodule is named as "name.0", "name.1", etc.
+    // Template parameter M must be a type derived from Module
     template <typename M, typename... Args>
     std::vector<std::shared_ptr<M>> register_modules(size_t layers, const std::string &name, Args &&...args) {
+        static_assert(std::is_base_of<Module, M>::value,
+                      "Template parameter M must be derived from infinicore::nn::Module");
+
         auto submodules = std::vector<std::shared_ptr<M>>(layers);
         for (size_t i = 0; i < layers; i++) {
-            register_module<M>(name + "." + std::to_string(i), std::forward<Args>(args)...);
+            submodules[i] = register_module<M>(name + "." + std::to_string(i), std::forward<Args>(args)...);
         }
         return submodules;
     }
 
 private:
-    void collect_all_parameters(const std::string &prefix, std::unordered_map<std::string, Parameter> &all_params) const;
+    void collect_all_parameters(std::unordered_map<std::string, Parameter> &all_params) const;
+    void update_submodule_paths();
 
 protected:
+    std::string path_;  // Full path of this module in the hierarchy (e.g., "layer1.sublayer")
     Device device_;
     std::unordered_map<std::string, std::shared_ptr<Module>> submodules_;
     std::unordered_map<std::string, Parameter> parameters_;
 };
+
+// ============================================================================
+// PyTorch-like Macros for Convenient Module Registration
+// ============================================================================
+
+/**
+ * @brief Register a submodule with automatic name inference from variable name
+ *
+ * Usage:
+ * @code
+ *   class MyModel : public Module {
+ *       INFINI_MODULE(Linear, layer1);
+ *       INFINI_MODULE(Linear, layer2);
+ *
+ *   public:
+ *       MyModel() {
+ *           INFINI_INIT_MODULE(layer1, 128, 64);
+ *           INFINI_INIT_MODULE(layer2, 64, 32);
+ *       }
+ *   };
+ * @endcode
+ */
+
+// Declare a module member variable
+#define INFINI_MODULE(ModuleType, name) \
+    std::shared_ptr<ModuleType> name##_
+
+// Initialize a module (use in constructor)
+#define INFINI_INIT_MODULE(name, ...) \
+    name##_ = register_module<std::remove_reference<decltype(*name##_)>::type>(#name, ##__VA_ARGS__)
+
+// Alternative: Combined declaration and initialization helper
+// Useful when you want to initialize inline (though less flexible)
+#define INFINI_SUBMODULE(ModuleType, name, ...) \
+    name##_ = register_module<ModuleType>(#name, ##__VA_ARGS__)
+
 } // namespace infinicore::nn

src/infinicore-test/main.cc

@@ -208,13 +208,29 @@ int main(int argc, char *argv[]) {
         auto results = runner.runAllTests();
         spdlog::debug("All tests completed");
 
-        // Count results
+        // Count results and collect failed tests
         size_t passed = 0, failed = 0;
+        std::vector<infinicore::test::TestResult> failed_tests;
         for (const auto &result : results) {
             if (result.passed) {
                 passed++;
             } else {
                 failed++;
+                failed_tests.push_back(result);
+            }
+        }
+
+        // Print list of failed tests if any
+        if (!failed_tests.empty()) {
+            std::cout << "\n==============================================\n"
+                      << "❌ FAILED TESTS\n"
+                      << "==============================================" << std::endl;
+            for (const auto &test : failed_tests) {
+                std::cout << "  • " << test.test_name;
+                if (!test.error_message.empty()) {
+                    std::cout << "\n    Error: " << test.error_message;
+                }
+                std::cout << "\n    Duration: " << test.duration.count() << "μs" << std::endl;
             }
         }
 
@@ -229,7 +245,7 @@ int main(int argc, char *argv[]) {
 
         // Exit with appropriate code
         if (failed > 0) {
-            std::cout << "\n❌ Some tests failed. Please review the output above." << std::endl;
+            std::cout << "\n❌ Some tests failed. Please review the failed tests list above." << std::endl;
             return EXIT_FAILURE;
         } else {
             std::cout << "\n✅ All tests passed!" << std::endl;

src/infinicore-test/memory_test.h

@@ -2,72 +2,17 @@
 #define __INFINICORE_MEMORY_TEST_H__
 
 #include "../infinicore/context/allocators/memory_allocator.hpp"
+#include "test_runner.h"
 #include <atomic>
 #include <cassert>
-#include <chrono>
-#include <exception>
 #include <future>
-#include <infinicore.hpp>
-#include <iostream>
-#include <memory>
 #include <mutex>
 #include <queue>
-#include <spdlog/spdlog.h>
 #include <thread>
 #include <unordered_map>
-#include <vector>
 
 namespace infinicore::test {
 
-// Test result structure
-struct TestResult {
-    std::string test_name;
-    bool passed;
-    std::string error_message;
-    std::chrono::microseconds duration;
-
-    TestResult(const std::string &name, bool pass, const std::string &error = "",
-               std::chrono::microseconds dur = std::chrono::microseconds(0))
-        : test_name(name), passed(pass), error_message(error), duration(dur) {}
-};
-
-// Test framework base class
-class MemoryTestFramework {
-public:
-    virtual ~MemoryTestFramework() = default;
-    virtual TestResult run() = 0;
-    virtual std::string getName() const = 0;
-
-protected:
-    void logTestStart(const std::string &test_name) {
-        std::cout << "[TEST] Starting: " << test_name << std::endl;
-    }
-
-    void logTestResult(const TestResult &result) {
-        std::cout << "[TEST] " << (result.passed ? "PASSED" : "FAILED")
-                  << ": " << result.test_name;
-        if (!result.passed && !result.error_message.empty()) {
-            std::cout << " - " << result.error_message;
-        }
-        std::cout << " (Duration: " << result.duration.count() << "μs)" << std::endl;
-    }
-
-    template <typename Func>
-    TestResult measureTime(const std::string &test_name, Func &&func) {
-        auto start = std::chrono::high_resolution_clock::now();
-        try {
-            bool result = func();
-            auto end = std::chrono::high_resolution_clock::now();
-            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
-            return TestResult(test_name, result, "", duration);
-        } catch (const std::exception &e) {
-            auto end = std::chrono::high_resolution_clock::now();
-            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
-            return TestResult(test_name, false, e.what(), duration);
-        }
-    }
-};
-
 // Mock allocator for testing exception safety
 class MockAllocator : public infinicore::MemoryAllocator {
 public:
@@ -149,13 +94,13 @@ class MemoryLeakDetector {
 };
 
 // Test categories
-class BasicMemoryTest : public MemoryTestFramework {
+class BasicMemoryTest : public TestFramework {
 public:
     TestResult run() override;
     std::string getName() const override { return "BasicMemoryTest"; }
 };
 
-class ConcurrencyTest : public MemoryTestFramework {
+class ConcurrencyTest : public TestFramework {
 public:
     TestResult run() override;
     std::string getName() const override { return "ConcurrencyTest"; }
@@ -166,7 +111,7 @@ class ConcurrencyTest : public MemoryTestFramework {
     TestResult testMemoryAllocationRace();
 };
 
-class ExceptionSafetyTest : public MemoryTestFramework {
+class ExceptionSafetyTest : public TestFramework {
 public:
     TestResult run() override;
     std::string getName() const override { return "ExceptionSafetyTest"; }
@@ -177,7 +122,7 @@ class ExceptionSafetyTest : public MemoryTestFramework {
     TestResult testContextSwitchException();
 };
 
-class MemoryLeakTest : public MemoryTestFramework {
+class MemoryLeakTest : public TestFramework {
 public:
     TestResult run() override;
     std::string getName() const override { return "MemoryLeakTest"; }
@@ -188,7 +133,7 @@ class MemoryLeakTest : public MemoryTestFramework {
     TestResult testExceptionLeakDetection();
 };
 
-class PerformanceTest : public MemoryTestFramework {
+class PerformanceTest : public TestFramework {
 public:
     TestResult run() override;
     std::string getName() const override { return "PerformanceTest"; }
@@ -199,7 +144,7 @@ class PerformanceTest : public MemoryTestFramework {
     TestResult testMemoryCopyPerformance();
 };
 
-class StressTest : public MemoryTestFramework {
+class StressTest : public TestFramework {
 public:
     TestResult run() override;
     std::string getName() const override { return "StressTest"; }