p70865324
/
InfiniTensor
forked from jiuyuan/InfiniTensor
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

allocate workspace from allocator for kunlun runtime

This commit is contained in:
kilinchange 2024-04-08 15:48:06 +08:00
parent e4387904c2
commit 3b7b5740af
2 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
include/kunlun/kunlun_runtime.h
View File

@ -22,9 +22,9 @@ class KUNLUNRuntimeObj : public RuntimeObj {
// 10GB for Longformer
// size_t longformerNum = 3lu * (1 << 30);
size_t workspaceSize = 3llu << 30; // 3 GB
KUNLUNPtr wkspacePtr = alloc(workspaceSize);
workspace =
make_ref<WorkspaceObj<KUNLUNPtr>>(wkspacePtr, workspaceSize);
// KUNLUNPtr wkspacePtr = alloc(workspaceSize);
// workspace =
// make_ref<WorkspaceObj<KUNLUNPtr>>(wkspacePtr, workspaceSize);
}
virtual ~KUNLUNRuntimeObj() {
KUNLUNPtr wkspacePtr = workspace->getWorkspace();
@ -33,6 +33,10 @@ class KUNLUNRuntimeObj : public RuntimeObj {
}
string toString() const override;
void init_workspace(KUNLUNPtr workspacePtr) {
workspace = make_ref<WorkspaceObj<KUNLUNPtr>>(wkspacePtr, workspaceSize);
}
void run(const Graph &graph, bool tune = false,
bool profiling = false) const;
// double runEvaluation(const Graph &graph, int nWarmups,

5
src/core/graph.cc
View File

@ -265,6 +265,11 @@ void GraphObj::dataMalloc(bool useNaiveAllocator, size_t memPoolSize) {
tensorToOffset[tensor.get()]));
}
}
if (runtime->isKUNLUN()) {
allocator.heapAlloc(runtime->workspaceSize);
runtime->init_workspace(allocator.getHeapPtr());
}
}
Tensor GraphObj::cloneKV(Tensor &tensor) {