diff --git a/.idea/.gitignore b/.idea/.gitignore
new file mode 100644
index 00000000..ab1f4164
--- /dev/null
+++ b/.idea/.gitignore
@@ -0,0 +1,10 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Ignored default folder with query files
+/queries/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml
+# Editor-based HTTP Client requests
+/httpRequests/
diff --git a/.idea/TinyInfiniTensor.iml b/.idea/TinyInfiniTensor.iml
new file mode 100644
index 00000000..d6ebd480
--- /dev/null
+++ b/.idea/TinyInfiniTensor.iml
@@ -0,0 +1,9 @@
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/misc.xml b/.idea/misc.xml
new file mode 100644
index 00000000..07115cdf
--- /dev/null
+++ b/.idea/misc.xml
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/modules.xml b/.idea/modules.xml
new file mode 100644
index 00000000..62de32f5
--- /dev/null
+++ b/.idea/modules.xml
@@ -0,0 +1,8 @@
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
new file mode 100644
index 00000000..67a1dcf4
--- /dev/null
+++ b/.idea/vcs.xml
@@ -0,0 +1,8 @@
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/include/core/allocator.h b/include/core/allocator.h
index 002601d2..5867c79e 100644
--- a/include/core/allocator.h
+++ b/include/core/allocator.h
@@ -27,7 +27,7 @@ namespace infini {
// TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并
// HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
// =================================== 作业 ===================================
-
+ std::map freeBlocksMap;
public:
Allocator(Runtime runtime);
diff --git a/include/core/graph.h b/include/core/graph.h
index c45580c2..4c534c8c 100644
--- a/include/core/graph.h
+++ b/include/core/graph.h
@@ -50,6 +50,7 @@ namespace infini
* so the topological sorting fails.
*/
bool topo_sort();
+ void reconstruct(Operator &op1, Operator &op2, Operator &op3);
void optimize();
diff --git a/src/core/allocator.cc b/src/core/allocator.cc
index ff593aef..80ff93c0 100644
--- a/src/core/allocator.cc
+++ b/src/core/allocator.cc
@@ -32,8 +32,31 @@ namespace infini
// =================================== 作业 ===================================
// TODO: 设计一个算法来分配内存，返回起始地址偏移量
// =================================== 作业 ===================================
+ used += size;
+ if(used > peak)
+ peak = used;
- return 0;
+ if(freeBlocksMap.empty()){
+ return peak - size;
+ }
+
+ bool freeBlkOk = false; //是否存在足够大（>= size）的空闲内存块
+ auto it = freeBlocksMap.begin();
+ for(; it != freeBlocksMap.end(); it++){
+ if(it->second >= size){
+ freeBlkOk = true;
+ if(it->second > size)
+ freeBlocksMap.insert({it->first + size, it->second - size});
+ break;
+ }
+ }
+ if(!freeBlkOk){
+ it--;
+ size_t moreMemory = size - it->second;
+ peak += moreMemory;
+ }
+ freeBlocksMap.erase(it);
+ return it->first;
}
void Allocator::free(size_t addr, size_t size)
@@ -44,6 +67,30 @@ namespace infini
// =================================== 作业 ===================================
// TODO: 设计一个算法来回收内存
// =================================== 作业 ===================================
+ auto it = freeBlocksMap.begin();
+ int flag = 0;
+ for(; it != freeBlocksMap.end(); it++){ //遍历空闲内存块，检查是否能与待回收内存合并成一个大内存块
+ if(it->first + it->second == addr){ //可以与待回收内存合并，该空闲块在前，待回收内存在后
+ flag = -1;
+ break;
+ }
+ else if(it->first == addr + size){ //可以与待回收内存合并，待回收内存在前，该空闲块在后，
+ flag = 1;
+ break;
+ }
+ }
+ if(flag == -1){
+ freeBlocksMap.insert({it->first, it->second + size}); //因为可以合并，插入合并后的大内存块，删去原来内存块
+ freeBlocksMap.erase(it);
+ }
+ else if(flag == 0){ //标记为0，说明以上遍历时没有找到能合并的空闲块
+ freeBlocksMap.insert({addr, size});
+ }
+ else if(flag == 1){
+ freeBlocksMap.insert({addr, size + it->second});
+ freeBlocksMap.erase(it);
+ }
+ used -= size;
}
void *Allocator::getPtr()
diff --git a/src/core/graph.cc b/src/core/graph.cc
index 3a906370..92e0992d 100644
--- a/src/core/graph.cc
+++ b/src/core/graph.cc
@@ -2,10 +2,13 @@
#include
#include
#include
+#include "operators/transpose.h"
+#include "operators/matmul.h"
+#include
+#include
namespace infini
{
-
void GraphObj::addOperatorAndConnect(const Operator &op)
{
sorted = false;
@@ -98,15 +101,128 @@ namespace infini
return this->sorted = true;
}
- void GraphObj::optimize()
- {
+ bool transposeOpsCancel(vector a, vector b){
+ if(a.size() != b.size())
+ return false;
+ for(int i = 0; i < (int)a.size(); i++){
+ if(b[a[i]] != i)
+ return false;
+ }
+ return true;
+ }
+
+ bool transOpCanIntegrateToMatmul(vector perm){
+ int size = perm.size();
+ for(int i = 0; i < size - 2; i++){
+ if(perm[i] != i)
+ return false;
+ }
+ if(perm[size - 2] != size - 1 || perm[size - 1] != size - 2)
+ return false;
+ else
+ return true;
+ }
+
+ void GraphObj::reconstruct(Operator &op1, Operator &op2, Operator &op3){
+ Tensor input = op1->getInputs(0);
+ if(input) {input->addTarget(op3);}
+ if(op2 == nullptr){
+ op3->replaceInput(op1->getOutput(), input);
+ op1->removeSuccessors(op3);
+ op3->removePredecessors(op1);
+ op1->getOutput()->removeTarget(op3);
+ }
+ else{
+ op3->replaceInput(op2->getOutput(), input);
+ op3->removePredecessors(op2);
+ op1->removeSuccessors(op2);
+ op1->getOutput()->removeTarget(op2);
+ }
+ for(auto &pred: op1->getPredecessors()){
+ op3->addPredecessors(pred);
+ pred->addSuccessors(op3);
+ }
+ }
+
+void GraphObj::optimize()
+{
// =================================== 作业 ===================================
// TODO: 设计一个算法来实现指定的图优化规则
// 图优化规则如下：
// 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除）
// 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去）
// =================================== 作业 ===================================
+ std::unordered_set toDelete; //待删除的算子
+ std::shared_ptr emptyPtr; //有的情况下重构计算图时需要
+ bool modified = true;
+ while(modified){
+ modified = false;
+ for(auto &op: ops){
+ if(toDelete.find(op.get()) != toDelete.end()){ //如果遍历到之前发现的待删除算子，略过
+ continue;
+ }
+ OpType opType = op->getOpType();
+ if(opType == OpType::Transpose){
+ TransposeObj* transOp = dynamic_cast(op.get());
+ Tensor input = op->getInputs(0);
+ for(auto &succ: op->getSuccessors()){
+ if(succ->getOpType() == OpType::Transpose){ //优化情况1 去除冗余transpose算子:当两个相邻算子都是transpose且相反操作
+ TransposeObj* transSucc = dynamic_cast(succ.get());
+ bool cancelOut = false; //判断两个相邻transpose算子是否相反的操作，在transposeOpsCancel函数中进行
+ if(transOp && transSucc){
+ cancelOut = transposeOpsCancel(transOp->getPermute(), transSucc->getPermute());
+ }
+ if(cancelOut){
+ toDelete.insert(succ.get());
+ removeTensor(succ->getOutput());
+ for(auto &succ_succ: succ->getSuccessors()){
+ reconstruct(op, succ, succ_succ);
+ }
+ if(op->getSuccessors().size() == 0){ //如果第1个transpose只有一个后继算子(即第2个transpose),那么可以删除第1个transpose
+ toDelete.insert(op.get());
+ //删除第1个transpose时需要进行以下重构
+ removeTensor(op->getOutput());
+ if(input) {input->removeTarget(op);}
+ for(auto &pred: op->getPredecessors())
+ pred->removeSuccessors(op);
+ }
+ modified = true;
+ }
+ }
+ else if(succ->getOpType() == OpType::MatMul){ //优化情况2 合并transpose与MatMul:当MatMul算子中含有属性transA、transB，且transpose对最后两个维度做交换
+ bool ok = transOpCanIntegrateToMatmul(transOp->getPermute()); //判断能否合并
+ if(ok){ //即使能够合并，也不一定就能删除transpose，除非transpose只有matMul唯一一个后继结点
+ MatmulObj* matmulSucc = dynamic_cast(succ.get());
+ if (matmulSucc){
+ if(succ->getInputs(0)->getGuid() == op->getOutput()->getGuid())
+ matmulSucc->setTransA(!(matmulSucc->getTransA()));
+ else
+ matmulSucc->setTransB(!(matmulSucc->getTransB()));
+ }
+ reconstruct(op, emptyPtr, succ); //合并transpose与matMul时需要进行重构：将transpose的输入作为matMul的输入
+ if(op->getSuccessors().size() == 0){ //如果transpose只有一个后继算子(即matMul)那么可以删除transpose。 size()==0是因为前面一步重构时已删去了两算子间的边
+ toDelete.insert(op.get());
+ removeTensor(op->getOutput());
+ if(input) {input->removeTarget(op);}
+ for(auto &pred: op->getPredecessors())
+ pred->removeSuccessors(op);
+ }
+ modified = true;
+ }
+ }
+ }
+ }
+ }
+ }
+ //删除之前遍历计算图时发现的冗余算子
+ for(int i = 0; i < (int)ops.size();){
+ if(toDelete.find(ops[i].get()) != toDelete.end())
+ ops.erase(ops.begin() + i);
+ else
+ i++;
}
+ this->sorted = false;
+}
Tensor GraphObj::getTensor(int fuid) const
{
@@ -152,10 +268,50 @@ namespace infini
// TODO：利用 allocator 给计算图分配内存
// HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存
// =================================== 作业 ===================================
+ std::unordered_map tensorOffset; //tensor在内存中的偏移量。 tensor id作为key,偏移量作为value
+ std::unordered_map tensorRefNum; //tensor被引用次数，即被多少个算子使用。tensor id作为key,被引用次数作为value
+ //记录每个张量的被引用次数
+ for(auto &tensor: tensors){
+ tensorRefNum[tensor->getFuid()] = static_cast(tensor->getTargets().size());
+ }
+ //为input分配内存
+ for(auto &g_input : getInputs()){
+ tensorOffset[g_input->getFuid()] = allocator.alloc(g_input->getBytes()); //if(!tensor->getSource())
+ }
+ //遍历算子
+ for(auto &op: ops){
+ for(auto &output: op->getOutputs()){
+ int id = output->getFuid();
+ if(tensorOffset.find(id) == tensorOffset.end()) //确保不会重复分配
+ tensorOffset[id] = allocator.alloc(output->getBytes());
+ }
+ for(auto &input: op->getInputs()){
+ int id = input->getFuid();
+ if(tensorRefNum.find(id) != tensorRefNum.end()){
+ tensorRefNum[id]--;
+ if(tensorRefNum[id] == 0){ //如果一个张量不再被使用((被引用次数是0)，可以释放其内存
+ auto it = tensorOffset.find(id); //释放前找到它在内存中的地址（偏移量）
+ if(it != tensorOffset.end()){
+ allocator.free(tensorOffset[id], input->getBytes());
+ //tensorOffset.erase(it);
+ }
+ tensorRefNum.erase(id);
+ }
+ }
+ }
+ }
+ void *memPtr = allocator.getPtr();
+ for(auto &tensor: tensors){
+ auto it = tensorOffset.find(tensor->getFuid());
+ IT_ASSERT(it != tensorOffset.end());
+ void *tensorPtr = static_cast(static_cast(memPtr) + it->second);
+ tensor->setDataBlob(make_ref(runtime, tensorPtr));
+ }
allocator.info();
}
+
Tensor GraphObj::addTensor(Shape dim, DataType dtype)
{
return tensors.emplace_back(make_ref(dim, dtype, runtime));
diff --git a/src/operators/concat.cc b/src/operators/concat.cc
index d1963308..7408414f 100644
--- a/src/operators/concat.cc
+++ b/src/operators/concat.cc
@@ -17,7 +17,10 @@ optional> ConcatObj::inferShape(const TensorVec &inputs) {
// TODO：修改 dims，返回正确的 concat 后的 shape
// REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13
// =================================== 作业 ===================================
-
+ int n = inputs.size();
+ for(int i = 1; i < n; i++){
+ dims[dim] += inputs[i]->getDims()[dim];
+ }
return {{dims}};
}
diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc
index 7a16ca27..b57063a8 100644
--- a/src/operators/matmul.cc
+++ b/src/operators/matmul.cc
@@ -27,7 +27,27 @@ namespace infini
// TODO：返回经过 matmul 操作后的 shape
// REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm
// =================================== 作业 ===================================
- return std::nullopt;
+ const auto A = inputs[0];
+ auto a_dim = A->getDims();
+ const auto B = inputs[1];
+ auto b_dim = B->getDims();
+ int size = a_dim.size();
+ Shape res = a_dim;
+ //multi-dimensional broadcasting except for the last two dimensions
+ for(int i = 0; i < size - 2; i++){
+ if(a_dim[i] < b_dim[i])
+ res[i] = b_dim[i];
+ }
+
+ if(transA)
+ res[size - 2] = a_dim[size - 1];
+ if(transB)
+ res[size - 1] = b_dim[size - 2];
+ else
+ res[size - 1] = b_dim[size - 1];
+
+ std::vector vec {res};
+ return vec;
}
} // namespace infini
\ No newline at end of file
diff --git a/src/operators/transpose.cc b/src/operators/transpose.cc
index faab2b69..31a009f1 100644
--- a/src/operators/transpose.cc
+++ b/src/operators/transpose.cc
@@ -33,8 +33,11 @@ namespace infini
// TODO：修改 output_dim，返回正确的 transpose 后的 shape
// REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21
// =================================== 作业 ===================================
-
- return std::nullopt;
+ for (int i = 0; i < rank; i++){
+ output_dim[i] = input_dim[transposePermute[i]];
+ }
+ std::vector res {output_dim};
+ return res;
}
std::string TransposeObj::toString() const
diff --git a/src/operators/unary.cc b/src/operators/unary.cc
index 3daad361..c7a8c928 100644
--- a/src/operators/unary.cc
+++ b/src/operators/unary.cc
@@ -39,7 +39,8 @@ namespace infini
// TODO：返回经过 clip 操作后的 shape
// REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13
// =================================== 作业 ===================================
- return std::nullopt;
+ const auto A = inputs[0];
+ return {{A->getDims()}};
}
std::string ClipObj::toString() const
@@ -66,7 +67,10 @@ namespace infini
// REF_FILE: src/core/operator.cc
// REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21
// =================================== 作业 ===================================
- return {};
+ DataType outputType = getOutputDataType();
+ int num = inputs.size();
+ std::vector res(num, outputType);
+ return res;
}
optional> CastObj::inferShape(const TensorVec &inputs)
@@ -75,7 +79,8 @@ namespace infini
// TODO：返回经过 cast 操作后的 shape
// REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21
// =================================== 作业 ===================================
- return std::nullopt;
+ const auto A = inputs[0];
+ return {{A->getDims()}};
}
std::string CastObj::toString() const
diff --git a/src/utils/operator_utils.cc b/src/utils/operator_utils.cc
index edbd2c82..37eab0d8 100644
--- a/src/utils/operator_utils.cc
+++ b/src/utils/operator_utils.cc
@@ -9,8 +9,32 @@ Shape infer_broadcast(const Shape &A, const Shape &B) {
// TODO：对 A 和 B 进行双向广播，返回广播后的形状。
// REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md
// =================================== 作业 ===================================
-
- return {};
+ Shape res;
+ long unsigned int i = 0;
+ if(A.size() > B.size()){
+ long unsigned int diff = A.size() - B.size();
+ for(; i < diff; i++)
+ res.push_back(A[i]);
+ for(; i < A.size(); i++){
+ if(A[i] == 1)
+ res.push_back(B[i-diff]);
+ else
+ res.push_back(A[i]);
+ }
+ }
+ else{
+ long unsigned int diff = B.size() - A.size();
+ for(; i < diff; i++)
+ res.push_back(B[i]);
+ for(; i < B.size(); i++){
+ if(B[i] == 1)
+ res.push_back(A[i-diff]);
+ else
+ res.push_back(B[i]);
+ }
+ }
+
+ return res;
}
int get_real_axis(const int &axis, const int &rank) {
diff --git a/test/kernels/nativecpu/test_nativecpu_concat.cc b/test/kernels/nativecpu/test_nativecpu_concat.cc
index fc87fb19..c1b0fa81 100644
--- a/test/kernels/nativecpu/test_nativecpu_concat.cc
+++ b/test/kernels/nativecpu/test_nativecpu_concat.cc
@@ -18,7 +18,6 @@ TEST(Concat, NativeCpu) {
t1->setData(IncrementalGenerator());
t2->setData(OneGenerator());
t3->setData(OneGenerator());
-
runtime->run(g);
EXPECT_TRUE(op->getOutput()->equalData(
vector{0, 1, 2, 1, 1, 1, 3, 4, 5, 1, 1, 1,