Complete TinyInfiniTensor training camp homework

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "3rd-party/nlohmann_json_cmake_fetchcontent"]
 	path = 3rd-party/nlohmann_json_cmake_fetchcontent
-	url = git@github.com:ArthurSonzogni/nlohmann_json_cmake_fetchcontent.git
+	url = https://github.com/ArthurSonzogni/nlohmann_json_cmake_fetchcontent.git
 [submodule "3rd-party/googletest"]
 	path = 3rd-party/googletest
-	url = git@github.com:google/googletest.git
+	url = https://github.com/google/googletest.git

CMakeLists.txt

@@ -29,7 +29,11 @@ endif()
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_EXTENSIONS OFF) # -std=gnu++11 when on, -std=c++11 when off
-add_compile_options(-Wno-error=unused-variable)
+if(MSVC)
+  add_compile_options(/W4 /utf-8 /Zc:preprocessor)
+else()
+  add_compile_options(-Wno-error=unused-variable)
+endif()
 
 find_package(
   Python
@@ -58,9 +62,14 @@ if(BUILD_TEST)
   include_directories(3rd-party/googletest/googletest/include)
 endif()
 
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -Wall -Werror -Wno-error=deprecated-declarations -Wno-error=pointer-arith")
-set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -UNDEBUG") # Enable assertion
-set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -UNDEBUG") # Enable assertion
+if(MSVC)
+  set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /UNDEBUG") # Enable assertion
+  set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} /UNDEBUG") # Enable assertion
+else()
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -Wall -Werror -Wno-error=deprecated-declarations -Wno-error=pointer-arith")
+  set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -UNDEBUG") # Enable assertion
+  set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -UNDEBUG") # Enable assertion
+endif()
 
 
 # Source files
@@ -72,7 +81,13 @@ if(USE_INTELCPU)
 endif()
 
 # Libraries
-add_library(InfiniTensor SHARED ${SRC})
+if(WIN32)
+  # The starter headers do not define DLL import/export annotations for static
+  # data members such as DataType::Float32, so use a static test library.
+  add_library(InfiniTensor STATIC ${SRC})
+else()
+  add_library(InfiniTensor SHARED ${SRC})
+endif()
 
 function(build_test files)
   # Non-recursive glob for skip failed tests
@@ -81,6 +96,12 @@ function(build_test files)
     get_filename_component(testname ${testsourcefile} NAME_WE)
     add_executable(${testname} ${testsourcefile})
     target_link_libraries(${testname} InfiniTensor GTest::gtest_main)
+    if(MSVC)
+      # Kernel registration is performed by static initializers in otherwise
+      # unreferenced translation units, so retain every object from the library.
+      target_link_options(${testname} PRIVATE
+                          "/WHOLEARCHIVE:$<TARGET_FILE:InfiniTensor>")
+    endif()
     add_test(NAME ${testname} COMMAND ${testname})
   endforeach(testsourcefile ${TEST_SOURCES})
 endfunction()

include/core/allocator.h

@@ -18,15 +18,18 @@ namespace infini {
 
     size_t peak;
 
+    // End of the currently occupied address range. Unlike peak, this may move
+    // backwards when the last block is released.
+    size_t currentEnd;
+
     size_t alignment;
 
     // pointer to the memory actually allocated
     void *ptr;
 
-    // =================================== 作业 ===================================
-    // TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并
-    // HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
-    // =================================== 作业 ===================================
+    // Free blocks indexed by their starting offsets. Keeping the map ordered
+    // makes it possible to coalesce adjacent blocks during free().
+    std::map<size_t, size_t> freeBlocks;
 
   public:
     Allocator(Runtime runtime);

src/core/allocator.cc

@@ -1,4 +1,5 @@
 #include "core/allocator.h"
+#include <algorithm>
 #include <utility>
 
 namespace infini
@@ -7,6 +8,7 @@ namespace infini
     {
         used = 0;
         peak = 0;
+        currentEnd = 0;
         ptr = nullptr;
 
         // 'alignment' defaults to sizeof(uint64_t), because it is the length of
@@ -26,38 +28,92 @@ namespace infini
     size_t Allocator::alloc(size_t size)
     {
         IT_ASSERT(this->ptr == nullptr);
-        // pad the size to the multiple of alignment
+        IT_ASSERT(size > 0);
         size = this->getAlignedSize(size);
 
-        // =================================== 作业 ===================================
-        // TODO: 设计一个算法来分配内存，返回起始地址偏移量
-        // =================================== 作业 ===================================
+        // First fit keeps allocation deterministic and reuses free space.
+        for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it)
+        {
+            if (it->second < size)
+                continue;
+
+            const size_t addr = it->first;
+            const size_t remaining = it->second - size;
+            freeBlocks.erase(it);
+            if (remaining > 0)
+                freeBlocks.emplace(addr + size, remaining);
+            used += size;
+            return addr;
+        }
 
-        return 0;
+        const size_t addr = currentEnd;
+        currentEnd += size;
+        peak = std::max(peak, currentEnd);
+        used += size;
+        return addr;
     }
 
     void Allocator::free(size_t addr, size_t size)
     {
         IT_ASSERT(this->ptr == nullptr);
+        IT_ASSERT(size > 0);
         size = getAlignedSize(size);
 
-        // =================================== 作业 ===================================
-        // TODO: 设计一个算法来回收内存
-        // =================================== 作业 ===================================
+        IT_ASSERT(addr % alignment == 0);
+        IT_ASSERT(addr <= currentEnd && size <= currentEnd - addr);
+        IT_ASSERT(size <= used);
+
+        auto next = freeBlocks.lower_bound(addr);
+        if (next != freeBlocks.end())
+            IT_ASSERT(addr + size <= next->first);
+        if (next != freeBlocks.begin())
+        {
+            const auto prev = std::prev(next);
+            IT_ASSERT(prev->first + prev->second <= addr);
+        }
+
+        auto current = freeBlocks.emplace_hint(next, addr, size);
+
+        next = std::next(current);
+        if (next != freeBlocks.end() &&
+            current->first + current->second == next->first)
+        {
+            current->second += next->second;
+            freeBlocks.erase(next);
+        }
+        if (current != freeBlocks.begin())
+        {
+            auto prev = std::prev(current);
+            if (prev->first + prev->second == current->first)
+            {
+                prev->second += current->second;
+                freeBlocks.erase(current);
+                current = prev;
+            }
+        }
+
+        used -= size;
+
+        if (current->first + current->second == currentEnd)
+        {
+            currentEnd = current->first;
+            freeBlocks.erase(current);
+        }
     }
 
     void *Allocator::getPtr()
     {
         if (this->ptr == nullptr)
         {
             this->ptr = runtime->alloc(this->peak);
-            printf("Allocator really alloc: %p %lu bytes\n", this->ptr, peak);
+            printf("Allocator really alloc: %p %zu bytes\n", this->ptr, peak);
         }
         return this->ptr;
     }
 
     size_t Allocator::getAlignedSize(size_t size)
     {
+        IT_ASSERT(size > 0);
         return ((size - 1) / this->alignment + 1) * this->alignment;
     }