From 464e2e8b86682eb6a3e000b93a34ee890638de0f Mon Sep 17 00:00:00 2001 From: lgeln10 <88012830+lgeln10@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:44:24 +0800 Subject: [PATCH] Complete TinyInfiniTensor training camp homework --- .gitmodules | 4 +- CMakeLists.txt | 31 ++++- include/core/allocator.h | 11 +- src/core/allocator.cc | 74 ++++++++++-- src/core/graph.cc | 219 ++++++++++++++++++++++++++++++++++-- src/operators/concat.cc | 14 ++- src/operators/matmul.cc | 26 ++++- src/operators/transpose.cc | 16 ++- src/operators/unary.cc | 19 +--- src/utils/operator_utils.cc | 15 ++- test/core/test_allocator.cc | 22 ++++ test/core/test_graph.cc | 60 +++++++++- 12 files changed, 441 insertions(+), 70 deletions(-) diff --git a/.gitmodules b/.gitmodules index e856b946..4800e247 100644 --- a/.gitmodules +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 836a7e09..a2f872c0 100644 --- a/CMakeLists.txt +++ b/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:$") + endif() add_test(NAME ${testname} COMMAND ${testname}) endforeach(testsourcefile ${TEST_SOURCES}) endfunction() diff --git a/include/core/allocator.h b/include/core/allocator.h index 002601d2..a5efa6e8 100644 --- a/include/core/allocator.h +++ b/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 freeBlocks; public: Allocator(Runtime runtime); diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..394aeaa2 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -1,4 +1,5 @@ #include "core/allocator.h" +#include #include 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,24 +28,77 @@ 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() @@ -51,13 +106,14 @@ namespace infini 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; } diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..a4a1a938 100644 --- a/src/core/graph.cc +++ b/src/core/graph.cc @@ -1,7 +1,10 @@ #include "core/graph.h" +#include "operators/matmul.h" +#include "operators/transpose.h" #include #include #include +#include namespace infini { @@ -100,12 +103,171 @@ namespace infini void GraphObj::optimize() { - // =================================== 作业 =================================== - // TODO: 设计一个算法来实现指定的图优化规则 - // 图优化规则如下： - // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除） - // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去） - // =================================== 作业 =================================== + // Graph outputs cannot be replaced because callers observe their tensor + // objects. Internal tensors may be removed after their users are rewired. + std::unordered_set graphOutputs; + for (const auto &tensor : tensors) + if (tensor->getTargets().empty()) + graphOutputs.insert(tensor.get()); + + // Rebuild all bidirectional links after each rewrite. This avoids stale + // predecessor/successor edges and also preserves duplicate input uses. + const auto rebuildConnections = [this]() + { + for (const auto &tensor : tensors) + { + tensor->targets.clear(); + tensor->source.reset(); + } + for (const auto &op : ops) + { + op->predecessors.clear(); + op->successors.clear(); + for (const auto &output : op->outputs) + output->source = op; + } + for (const auto &op : ops) + { + for (const auto &input : op->inputs) + { + input->targets.emplace_back(op); + if (auto predecessor = input->source.lock()) + { + predecessor->successors.emplace_back(op); + op->predecessors.emplace_back(predecessor); + } + } + } + }; + + const auto eraseOperator = [this](const Operator &op) + { + ops.erase(std::remove(ops.begin(), ops.end(), op), ops.end()); + }; + const auto eraseTensor = [this](const Tensor &tensor) + { + tensors.erase(std::remove(tensors.begin(), tensors.end(), tensor), + tensors.end()); + }; + const auto areInverse = [](const vector &first, + const vector &second) + { + if (first.size() != second.size()) + return false; + for (size_t i = 0; i < first.size(); ++i) + { + if (second[i] < 0 || + second[i] >= static_cast(first.size()) || + first[second[i]] != static_cast(i)) + return false; + } + return true; + }; + const auto swapsLastTwoAxes = [](const vector &permute) + { + if (permute.size() < 2) + return false; + for (size_t i = 0; i + 2 < permute.size(); ++i) + if (permute[i] != static_cast(i)) + return false; + const size_t rank = permute.size(); + return permute[rank - 2] == static_cast(rank - 1) && + permute[rank - 1] == static_cast(rank - 2); + }; + + rebuildConnections(); + bool changed = true; + while (changed) + { + changed = false; + + // Rule 1: remove two inverse transposes. The tensor between them + // must be exclusively consumed by the second transpose. Otherwise + // removing the first transpose would break its other consumers. + for (const auto &secondOp : ops) + { + if (secondOp->getOpType() != OpType::Transpose) + continue; + + const Tensor middle = secondOp->inputs[0]; + const Operator firstOp = middle->source.lock(); + if (!firstOp || firstOp->getOpType() != OpType::Transpose) + continue; + + const auto middleTargets = middle->getTargets(); + if (middleTargets.size() != 1 || middleTargets[0] != secondOp) + continue; + + const auto first = as(firstOp); + const auto second = as(secondOp); + if (!areInverse(first->getPermute(), second->getPermute())) + continue; + + const Tensor result = secondOp->outputs[0]; + if (graphOutputs.count(result.get()) != 0) + continue; + + const Tensor original = firstOp->inputs[0]; + for (const auto &consumer : result->getTargets()) + consumer->replaceInput(result, original); + + eraseOperator(secondOp); + eraseOperator(firstOp); + eraseTensor(result); + eraseTensor(middle); + rebuildConnections(); + sorted = false; + changed = true; + break; + } + if (changed) + continue; + + // Rule 2: fold a last-two-axis transpose into MatMul. The transpose + // output must be exclusive to this MatMul; a shared tensor must keep + // its producer and is deliberately left untouched. + for (const auto &op : ops) + { + if (op->getOpType() != OpType::MatMul) + continue; + + auto matmul = as(op); + for (size_t inputIndex = 0; inputIndex < op->inputs.size(); + ++inputIndex) + { + const Tensor transposed = op->inputs[inputIndex]; + const Operator transposeOp = transposed->source.lock(); + if (!transposeOp || + transposeOp->getOpType() != OpType::Transpose) + continue; + + const auto targets = transposed->getTargets(); + if (targets.size() != 1 || targets[0] != op) + continue; + + auto transpose = as(transposeOp); + if (!swapsLastTwoAxes(transpose->getPermute())) + continue; + + op->inputs[inputIndex] = transposeOp->inputs[0]; + if (inputIndex == 0) + matmul->setTransA(!matmul->getTransA()); + else + matmul->setTransB(!matmul->getTransB()); + + eraseOperator(transposeOp); + eraseTensor(transposed); + rebuildConnections(); + sorted = false; + changed = true; + break; + } + if (changed) + break; + } + } + + IT_ASSERT(checkValid()); } Tensor GraphObj::getTensor(int fuid) const @@ -145,14 +307,47 @@ namespace infini void GraphObj::dataMalloc() { - // topological sorting first IT_ASSERT(topo_sort() == true); - // =================================== 作业 =================================== - // TODO：利用 allocator 给计算图分配内存 - // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存 - // =================================== 作业 =================================== + std::map offsets; + std::map remainingUses; + for (const auto &tensor : tensors) + remainingUses[tensor] = tensor->getTargets().size(); + + const auto allocateTensor = [this, &offsets](const Tensor &tensor) + { + if (offsets.find(tensor) == offsets.end()) + offsets[tensor] = allocator.alloc(tensor->getBytes()); + }; + + // Graph inputs are live before execution starts. + for (const auto &tensor : tensors) + if (!tensor->getSource()) + allocateTensor(tensor); + + // Allocate outputs before releasing the current inputs: kernels cannot + // overwrite an input while they are still producing the output. + for (const auto &op : ops) + { + for (const auto &output : op->getOutputs()) + allocateTensor(output); + for (const auto &input : op->getInputs()) + { + auto &uses = remainingUses[input]; + IT_ASSERT(uses > 0); + --uses; + if (uses == 0) + allocator.free(offsets.at(input), input->getBytes()); + } + } + + void *base = allocator.getPtr(); + for (const auto &tensor : tensors) + { + auto *data = static_cast(base) + offsets.at(tensor); + tensor->setDataBlob(make_ref(runtime, data)); + } allocator.info(); } @@ -227,4 +422,4 @@ namespace infini return true; } -} // namespace infini \ No newline at end of file +} // namespace infini diff --git a/src/operators/concat.cc b/src/operators/concat.cc index d1963308..71253cd7 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -10,13 +10,19 @@ ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim) } optional> ConcatObj::inferShape(const TensorVec &inputs) { + IT_ASSERT(!inputs.empty()); Shape dims = inputs[0]->getDims(); auto rank = inputs[0]->getRank(); - // =================================== 作业 =================================== - // TODO：修改 dims，返回正确的 concat 后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13 - // =================================== 作业 =================================== + for (size_t i = 1; i < inputs.size(); ++i) { + const auto inputDims = inputs[i]->getDims(); + IT_ASSERT(inputs[i]->getRank() == rank); + for (size_t axis = 0; axis < rank; ++axis) { + if (axis != static_cast(dim)) + IT_ASSERT(inputDims[axis] == dims[axis]); + } + dims[dim] += inputDims[dim]; + } return {{dims}}; } diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..3ebdb44b 100644 --- a/src/operators/matmul.cc +++ b/src/operators/matmul.cc @@ -1,4 +1,5 @@ #include "operators/matmul.h" +#include "utils/operator_utils.h" namespace infini { @@ -23,11 +24,24 @@ namespace infini optional> MatmulObj::inferShape(const TensorVec &inputs) { - // =================================== 作业 =================================== - // TODO：返回经过 matmul 操作后的 shape - // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm - // =================================== 作业 =================================== - return std::nullopt; + const Shape shapeA = inputs[0]->getDims(); + const Shape shapeB = inputs[1]->getDims(); + IT_ASSERT(shapeA.size() >= 2 && shapeB.size() >= 2); + + const size_t rankA = shapeA.size(); + const size_t rankB = shapeB.size(); + m = transA ? shapeA[rankA - 1] : shapeA[rankA - 2]; + k = transA ? shapeA[rankA - 2] : shapeA[rankA - 1]; + const int kB = transB ? shapeB[rankB - 1] : shapeB[rankB - 2]; + n = transB ? shapeB[rankB - 2] : shapeB[rankB - 1]; + IT_ASSERT(k == kB); + + const Shape batchA(shapeA.begin(), shapeA.end() - 2); + const Shape batchB(shapeB.begin(), shapeB.end() - 2); + Shape output = infer_broadcast(batchA, batchB); + output.push_back(m); + output.push_back(n); + return {{output}}; } -} // namespace infini \ No newline at end of file +} // namespace infini diff --git a/src/operators/transpose.cc b/src/operators/transpose.cc index faab2b69..ea473082 100644 --- a/src/operators/transpose.cc +++ b/src/operators/transpose.cc @@ -9,6 +9,7 @@ namespace infini auto rank = input->getRank(); if (permute.empty()) { + transposePermute.resize(rank); for (size_t i = 0; i < rank; ++i) { transposePermute[i] = i; @@ -29,12 +30,17 @@ namespace infini auto output_dim = input_dim; int rank = A->getRank(); - // =================================== 作业 =================================== - // TODO：修改 output_dim，返回正确的 transpose 后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21 - // =================================== 作业 =================================== + vector seen(rank, false); + for (int i = 0; i < rank; ++i) + { + const int axis = transposePermute[i]; + IT_ASSERT(axis >= 0 && axis < rank); + IT_ASSERT(!seen[axis]); + seen[axis] = true; + output_dim[i] = input_dim[axis]; + } - return std::nullopt; + return {{output_dim}}; } std::string TransposeObj::toString() const diff --git a/src/operators/unary.cc b/src/operators/unary.cc index 3daad361..efd9ec8f 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -35,11 +35,7 @@ namespace infini optional> ClipObj::inferShape(const TensorVec &inputs) { - // =================================== 作业 =================================== - // TODO：返回经过 clip 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 - // =================================== 作业 =================================== - return std::nullopt; + return {{inputs[0]->getDims()}}; } std::string ClipObj::toString() const @@ -61,21 +57,12 @@ namespace infini vector CastObj::inferDataType(const TensorVec &inputs) const { - // =================================== 作业 =================================== - // TODO：返回经过 cast 操作后, 输出 tensor 的数目和数据类型 - // REF_FILE: src/core/operator.cc - // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 - // =================================== 作业 =================================== - return {}; + return {getOutputDataType()}; } optional> CastObj::inferShape(const TensorVec &inputs) { - // =================================== 作业 =================================== - // TODO：返回经过 cast 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 - // =================================== 作业 =================================== - return std::nullopt; + return {{inputs[0]->getDims()}}; } std::string CastObj::toString() const diff --git a/src/utils/operator_utils.cc b/src/utils/operator_utils.cc index edbd2c82..8caab192 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -4,13 +4,16 @@ namespace infini { Shape infer_broadcast(const Shape &A, const Shape &B) { + const size_t rank = std::max(A.size(), B.size()); + Shape result(rank, 1); - // =================================== 作业 =================================== - // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 - // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md - // =================================== 作业 =================================== - - return {}; + for (size_t offset = 0; offset < rank; ++offset) { + const int a = offset < A.size() ? A[A.size() - 1 - offset] : 1; + const int b = offset < B.size() ? B[B.size() - 1 - offset] : 1; + IT_ASSERT(a == b || a == 1 || b == 1); + result[rank - 1 - offset] = std::max(a, b); + } + return result; } int get_real_axis(const int &axis, const int &rank) { diff --git a/test/core/test_allocator.cc b/test/core/test_allocator.cc index 0515edc9..971a10fc 100644 --- a/test/core/test_allocator.cc +++ b/test/core/test_allocator.cc @@ -71,4 +71,26 @@ namespace infini EXPECT_EQ(ptr1, ptr2); } + TEST(Allocator, graphDataMallocReusesExpiredTensorSafely) + { + Runtime runtime = NativeCpuRuntimeObj::getInstance(); + Graph g = make_ref(runtime); + Tensor input = g->addTensor({1, 2, 2, 3}, DataType::Float32); + auto first = g->addOp(input, nullptr); + auto second = g->addOp(first->getOutput(), nullptr); + + g->dataMalloc(); + + auto *inputPtr = input->getRawDataPtr(); + auto *middlePtr = first->getOutput()->getRawDataPtr(); + auto *outputPtr = second->getOutput()->getRawDataPtr(); + EXPECT_EQ(inputPtr, outputPtr); + EXPECT_NE(inputPtr, middlePtr); + + input->setData(IncrementalGenerator()); + runtime->run(g); + EXPECT_TRUE(second->getOutput()->equalData( + vector{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11})); + } + } // namespace infini diff --git a/test/core/test_graph.cc b/test/core/test_graph.cc index bf696ddb..e1a5c6ee 100644 --- a/test/core/test_graph.cc +++ b/test/core/test_graph.cc @@ -3,6 +3,7 @@ #include "core/runtime.h" #include "operators/matmul.h" #include "operators/transpose.h" +#include "operators/unary.h" #include "test.h" @@ -37,4 +38,61 @@ namespace infini EXPECT_EQ(op->getTransA(), false); EXPECT_EQ(op->getTransB(), true); } -} \ No newline at end of file + + TEST(Graph, DoNotRemoveInverseTransposesWhenIntermediateIsShared) + { + Runtime runtime = NativeCpuRuntimeObj::getInstance(); + Graph g = make_ref(runtime); + Tensor input = g->addTensor({2, 3}, DataType::Float32); + Tensor middle = g->addTensor({3, 2}, DataType::Float32); + Tensor restored = g->addTensor({2, 3}, DataType::Float32); + Tensor sharedOutput = g->addTensor({3, 2}, DataType::Float32); + Tensor restoredOutput = g->addTensor({2, 3}, DataType::Float32); + + auto first = g->addOpWithOutputs(input, middle, + Shape{1, 0}); + auto second = g->addOpWithOutputs(middle, restored, + Shape{1, 0}); + g->addOpWithOutputs(middle, sharedOutput, std::nullopt, + std::nullopt); + g->addOpWithOutputs(restored, restoredOutput, std::nullopt, + std::nullopt); + + g->optimize(); + + EXPECT_EQ(g->getOperators().size(), 4); + EXPECT_EQ(g->getTensors().size(), 5); + EXPECT_EQ(middle->getSource(), first); + EXPECT_EQ(restored->getSource(), second); + EXPECT_EQ(middle->getTargets().size(), 2); + EXPECT_TRUE(g->checkValid()); + } + + TEST(Graph, DoNotFoldTransposeIntoMatmulWhenOutputIsShared) + { + Runtime runtime = NativeCpuRuntimeObj::getInstance(); + Graph g = make_ref(runtime); + Tensor input = g->addTensor({2, 3}, DataType::Float32); + Tensor transposed = g->addTensor({3, 2}, DataType::Float32); + Tensor rhs = g->addTensor({2, 4}, DataType::Float32); + Tensor matmulOutput = g->addTensor({3, 4}, DataType::Float32); + Tensor sharedOutput = g->addTensor({3, 2}, DataType::Float32); + + auto transpose = g->addOpWithOutputs(input, transposed, + Shape{1, 0}); + auto matmul = g->addOpWithOutputs(transposed, rhs, + matmulOutput); + g->addOpWithOutputs(transposed, sharedOutput, std::nullopt, + std::nullopt); + + g->optimize(); + + EXPECT_EQ(g->getOperators().size(), 3); + EXPECT_EQ(g->getTensors().size(), 5); + EXPECT_EQ(transposed->getSource(), transpose); + EXPECT_EQ(transposed->getTargets().size(), 2); + EXPECT_EQ(matmul->getInputs(0), transposed); + EXPECT_FALSE(matmul->getTransA()); + EXPECT_TRUE(g->checkValid()); + } +}