diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..9230d014 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -1,69 +1,121 @@ -#include "core/allocator.h" -#include - -namespace infini -{ - Allocator::Allocator(Runtime runtime) : runtime(runtime) - { - used = 0; - peak = 0; - ptr = nullptr; - - // 'alignment' defaults to sizeof(uint64_t), because it is the length of - // the longest data type currently supported by the DataType field of - // the tensor - alignment = sizeof(uint64_t); - } - - Allocator::~Allocator() - { - if (this->ptr != nullptr) - { - runtime->dealloc(this->ptr); - } - } - - size_t Allocator::alloc(size_t size) - { - IT_ASSERT(this->ptr == nullptr); - // pad the size to the multiple of alignment - size = this->getAlignedSize(size); - - // =================================== 作业 =================================== - // TODO: 设计一个算法来分配内存，返回起始地址偏移量 - // =================================== 作业 =================================== - - return 0; - } - - void Allocator::free(size_t addr, size_t size) - { - IT_ASSERT(this->ptr == nullptr); - size = getAlignedSize(size); - - // =================================== 作业 =================================== - // TODO: 设计一个算法来回收内存 - // =================================== 作业 =================================== - } - - void *Allocator::getPtr() - { - if (this->ptr == nullptr) - { - this->ptr = runtime->alloc(this->peak); - printf("Allocator really alloc: %p %lu bytes\n", this->ptr, peak); - } - return this->ptr; - } - - size_t Allocator::getAlignedSize(size_t size) - { - return ((size - 1) / this->alignment + 1) * this->alignment; - } - - void Allocator::info() - { - std::cout << "Used memory: " << this->used - << ", peak memory: " << this->peak << std::endl; - } -} +#include "core/allocator.h" +#include + +namespace infini +{ + Allocator::Allocator(Runtime runtime) : runtime(runtime) + { + used = 0; + peak = 0; + ptr = nullptr; + + // 'alignment' defaults to sizeof(uint64_t), because it is the length of + // the longest data type currently supported by the DataType field of + // the tensor + alignment = sizeof(uint64_t); + } + + Allocator::~Allocator() + { + if (this->ptr != nullptr) + { + runtime->dealloc(this->ptr); + } + } + + size_t Allocator::alloc(size_t size) + { + IT_ASSERT(this->ptr == nullptr); + // pad the size to the multiple of alignment + size = this->getAlignedSize(size); + + // =================================== 作业 =================================== + // TODO: 设计一个算法来分配内存，返回起始地址偏移量 + // =================================== 作业 =================================== + // First-fit: search free blocks for a suitable one + for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it) { + if (it->second >= size) { + size_t addr = it->first; + size_t remaining = it->second - size; + freeBlocks.erase(it); + if (remaining > 0) { + freeBlocks[addr + size] = remaining; + } + return addr; + } + } + // No suitable free block found, allocate at the end + size_t addr = used; + used += size; + peak = std::max(peak, used); + return addr; + } + + void Allocator::free(size_t addr, size_t size) + { + IT_ASSERT(this->ptr == nullptr); + size = getAlignedSize(size); + + // =================================== 作业 =================================== + // TODO: 设计一个算法来回收内存 + // =================================== 作业 =================================== + // If the block is at the end of used memory, shrink used + if (addr + size == used) { + used = addr; + // Remove any trailing free blocks that are now past the new used + while (!freeBlocks.empty()) { + auto last = std::prev(freeBlocks.end()); + if (last->first + last->second == addr) { + addr = last->first; + size += last->second; + freeBlocks.erase(last); + } else { + break; + } + } + used = addr; + return; + } + + // Insert into free blocks + auto [it, inserted] = freeBlocks.insert({addr, size}); + IT_ASSERT(inserted); + + // Merge with next block if adjacent + auto next = std::next(it); + if (next != freeBlocks.end() && addr + size == next->first) { + it->second += next->second; + freeBlocks.erase(next); + } + + // Merge with previous block if adjacent + if (it != freeBlocks.begin()) { + auto prev = std::prev(it); + if (prev->first + prev->second == addr) { + prev->second += it->second; + freeBlocks.erase(it); + } + } + } + + void *Allocator::getPtr() + { + if (this->ptr == nullptr) + { + this->ptr = runtime->alloc(this->peak); + printf("Allocator really alloc: %p %lu bytes\n", this->ptr, peak); + } + return this->ptr; + } + + size_t Allocator::getAlignedSize(size_t size) + { + return ((size - 1) / this->alignment + 1) * this->alignment; + } + + void Allocator::info() + { + std::cout << "Used memory: " << this->used + << ", peak memory: " << this->peak << std::endl; + } +} diff --git a/src/core/data_type.cc b/src/core/data_type.cc index 3825c9cc..ccfdee5a 100644 --- a/src/core/data_type.cc +++ b/src/core/data_type.cc @@ -1,22 +1,22 @@ -#include "core/data_type.h" - -namespace infini { -// Move implementation here to avoid compile time error on some platform -// to be consistent with onnx -// https://github.com/onnx/onnx/blob/aeb21329122b96df1d3ef33b500a35ca140b1431/onnx/onnx.proto#L484 -const DataType DataType::Undefine(0); -const DataType DataType::Float32(1); -const DataType DataType::UInt8(2); -const DataType DataType::Int8(3); -const DataType DataType::UInt16(4); -const DataType DataType::Int16(5); -const DataType DataType::Int32(6); -const DataType DataType::Int64(7); -const DataType DataType::String(8); -const DataType DataType::Bool(9); -const DataType DataType::Float16(10); -const DataType DataType::Double(11); -const DataType DataType::UInt32(12); -const DataType DataType::UInt64(13); -const DataType DataType::BFloat16(16); -} // namespace infini +#include "core/data_type.h" + +namespace infini { +// Move implementation here to avoid compile time error on some platform +// to be consistent with onnx +// https://github.com/onnx/onnx/blob/aeb21329122b96df1d3ef33b500a35ca140b1431/onnx/onnx.proto#L484 +const DataType DataType::Undefine(0); +const DataType DataType::Float32(1); +const DataType DataType::UInt8(2); +const DataType DataType::Int8(3); +const DataType DataType::UInt16(4); +const DataType DataType::Int16(5); +const DataType DataType::Int32(6); +const DataType DataType::Int64(7); +const DataType DataType::String(8); +const DataType DataType::Bool(9); +const DataType DataType::Float16(10); +const DataType DataType::Double(11); +const DataType DataType::UInt32(12); +const DataType DataType::UInt64(13); +const DataType DataType::BFloat16(16); +} // namespace infini diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..b014aa43 100644 --- a/src/core/graph.cc +++ b/src/core/graph.cc @@ -1,230 +1,448 @@ -#include "core/graph.h" -#include -#include -#include - -namespace infini -{ - - void GraphObj::addOperatorAndConnect(const Operator &op) - { - sorted = false; - ops.push_back(op); - for (auto &input : op->getInputs()) - { - if (input) - { - input->addTarget(op); - if (auto pred = input->getSource()) - { - pred->addSuccessors(op); - op->addPredecessors(pred); - } - } - } - for (auto &output : op->getOutputs()) - { - if (output) - { - output->setSource(op); - for (auto &succ : output->getTargets()) - { - succ->addPredecessors(op); - op->addSuccessors(succ); - } - } - } - } - - string GraphObj::toString() const - { - std::ostringstream oss; - oss << "Graph Tensors:\n"; - for (const auto &tensor : tensors) - oss << tensor << "\n"; - - oss << "Graph operators:\n"; - for (const auto &op : ops) - { - vector preds, succs; - for (auto &o : op->getPredecessors()) - preds.emplace_back(o->getGuid()); - for (auto &o : op->getSuccessors()) - succs.emplace_back(o->getGuid()); - oss << "OP " << op->getGuid(); - oss << ", pred " << vecToString(preds); - oss << ", succ " << vecToString(succs); - oss << ", " << op << "\n"; - } - return oss.str(); - } - - bool GraphObj::topo_sort() - { - if (this->sorted) - { - return true; - } - std::vector sorted; - std::unordered_set flags; - sorted.reserve(ops.size()); - flags.reserve(ops.size()); - while (sorted.size() < ops.size()) - { - // Any node is move to sorted in this loop. - auto modified = false; - for (auto const &op : ops) - { - if (auto const &inputs = op->getInputs(); - flags.find(op.get()) == flags.end() && - std::all_of(inputs.begin(), inputs.end(), - [&flags](auto const &input) - { - auto ptr = input->getSource().get(); - return !ptr || flags.find(ptr) != flags.end(); - })) - { - modified = true; - sorted.emplace_back(op); - flags.insert(op.get()); - } - } - if (!modified) - { - return false; - } - } - this->ops = std::move(sorted); - return this->sorted = true; - } - - void GraphObj::optimize() - { - // =================================== 作业 =================================== - // TODO: 设计一个算法来实现指定的图优化规则 - // 图优化规则如下： - // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除） - // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去） - // =================================== 作业 =================================== - } - - Tensor GraphObj::getTensor(int fuid) const - { - for (auto tensor : tensors) - { - if (tensor->getFuid() == fuid) - { - return tensor; - } - } - return nullptr; - } - - void GraphObj::shape_infer() - { - for (auto &op : ops) - { - auto ans = op->inferShape(); - IT_ASSERT(ans.has_value()); - auto oldOutputs = op->getOutputs(); - IT_ASSERT(ans.value().size() == oldOutputs.size()); - // replace the old outputshape and size with new one - for (int i = 0; i < (int)ans.value().size(); ++i) - { - auto newShape = ans.value()[i]; - auto oldShape = oldOutputs[i]->getDims(); - auto fuid = oldOutputs[i]->getFuid(); - if (newShape != oldShape) - { - auto tensor = this->getTensor(fuid); - tensor->setShape(newShape); - } - } - } - } - - void GraphObj::dataMalloc() - { - // topological sorting first - IT_ASSERT(topo_sort() == true); - - // =================================== 作业 =================================== - // TODO：利用 allocator 给计算图分配内存 - // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存 - // =================================== 作业 =================================== - - allocator.info(); - } - - Tensor GraphObj::addTensor(Shape dim, DataType dtype) - { - return tensors.emplace_back(make_ref(dim, dtype, runtime)); - } - - Tensor GraphObj::addTensor(const Tensor &tensor) - { - IT_ASSERT(tensor->getRuntime() == runtime, - std::string("Tensor runtime mismatch: cannot add a tenosr in ") + - tensor->getRuntime()->toString() + " to " + - runtime->toString()); - tensors.emplace_back(tensor); - return tensor; - } - - TensorVec GraphObj::addTensor(const TensorVec &tensors) - { - for (auto &t : tensors) - addTensor(t); - return tensors; - } - - // tensor's "source" and "target" must be in "ops". - // tensor has no "source" and no "target" must not exist. - // "inputs" or "outputs" of operators must be in "tensors" - // "predecessors" and "successors" of an operator of "ops" must be in "ops". - bool GraphObj::checkValid() const - { - for (auto tensor : tensors) - { - IT_ASSERT(!(tensor->getTargets().size() == 0 && - nullptr == tensor->getSource())); - for (auto op : tensor->getTargets()) - { - IT_ASSERT(std::find(ops.begin(), ops.end(), op) != ops.end()); - } - auto op = tensor->getSource(); - IT_ASSERT(!(op && std::find(ops.begin(), ops.end(), op) == ops.end())); - } - for (auto op : ops) - { - for (auto tensor : op->getInputs()) - { - IT_ASSERT(std::find(tensors.begin(), tensors.end(), tensor) != - tensors.end()); - } - for (auto tensor : op->getOutputs()) - { - IT_ASSERT(std::find(tensors.begin(), tensors.end(), tensor) != - tensors.end()); - } - for (auto pre : op->getPredecessors()) - { - IT_ASSERT(std::find(ops.begin(), ops.end(), pre) != ops.end()); - } - for (auto suc : op->getSuccessors()) - { - IT_ASSERT(std::find(ops.begin(), ops.end(), suc) != ops.end()); - } - } - std::set s; - // check whether two tensors with the same FUID exist - for (auto tensor : tensors) - { - int cnt = s.count(tensor->getFuid()); - IT_ASSERT(cnt == 0, std::to_string(tensor->getFuid())); - s.insert(tensor->getFuid()); - } - return true; - } - +#include "core/graph.h" +#include "operators/matmul.h" +#include "operators/transpose.h" +#include +#include +#include + +namespace infini +{ + + void GraphObj::addOperatorAndConnect(const Operator &op) + { + sorted = false; + ops.push_back(op); + for (auto &input : op->getInputs()) + { + if (input) + { + input->addTarget(op); + if (auto pred = input->getSource()) + { + pred->addSuccessors(op); + op->addPredecessors(pred); + } + } + } + for (auto &output : op->getOutputs()) + { + if (output) + { + output->setSource(op); + for (auto &succ : output->getTargets()) + { + succ->addPredecessors(op); + op->addSuccessors(succ); + } + } + } + } + + string GraphObj::toString() const + { + std::ostringstream oss; + oss << "Graph Tensors:\n"; + for (const auto &tensor : tensors) + oss << tensor << "\n"; + + oss << "Graph operators:\n"; + for (const auto &op : ops) + { + vector preds, succs; + for (auto &o : op->getPredecessors()) + preds.emplace_back(o->getGuid()); + for (auto &o : op->getSuccessors()) + succs.emplace_back(o->getGuid()); + oss << "OP " << op->getGuid(); + oss << ", pred " << vecToString(preds); + oss << ", succ " << vecToString(succs); + oss << ", " << op << "\n"; + } + return oss.str(); + } + + bool GraphObj::topo_sort() + { + if (this->sorted) + { + return true; + } + std::vector sorted; + std::unordered_set flags; + sorted.reserve(ops.size()); + flags.reserve(ops.size()); + while (sorted.size() < ops.size()) + { + // Any node is move to sorted in this loop. + auto modified = false; + for (auto const &op : ops) + { + if (auto const &inputs = op->getInputs(); + flags.find(op.get()) == flags.end() && + std::all_of(inputs.begin(), inputs.end(), + [&flags](auto const &input) + { + auto ptr = input->getSource().get(); + return !ptr || flags.find(ptr) != flags.end(); + })) + { + modified = true; + sorted.emplace_back(op); + flags.insert(op.get()); + } + } + if (!modified) + { + return false; + } + } + this->ops = std::move(sorted); + return this->sorted = true; + } + + void GraphObj::optimize() + { + // =================================== 作业 =================================== + // TODO: 设计一个算法来实现指定的图优化规则 + // 图优化规则如下： + // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除） + // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去） + // =================================== 作业 =================================== + + bool changed = true; + while (changed) { + changed = false; + + // Rule 1: Remove redundant consecutive transposes + for (auto it = ops.begin(); it != ops.end(); ) { + auto op = *it; + if (op->getOpType() != OpType::Transpose) { + ++it; + continue; + } + auto transpose1 = as(op); + auto output = transpose1->getOutput(); + auto targets = output->getTargets(); + if (targets.size() != 1) { + ++it; + continue; + } + auto nextOp = targets[0]; + if (nextOp->getOpType() != OpType::Transpose) { + ++it; + continue; + } + auto transpose2 = as(nextOp); + const auto &perm1 = transpose1->getPermute(); + const auto &perm2 = transpose2->getPermute(); + // Check if perm1 and perm2 are inverses + bool areInverses = true; + for (size_t i = 0; i < perm1.size(); ++i) { + if (perm2[perm1[i]] != (int)i) { + areInverses = false; + break; + } + } + if (!areInverses) { + ++it; + continue; + } + // Remove both transposes: connect input of transpose1 to consumers of transpose2 + auto input = transpose1->getInputs(0); + auto output2 = transpose2->getOutput(); + + // Replace output2 with input in all consumers of output2 + auto consumers = output2->getTargets(); + for (auto &consumer : consumers) { + consumer->replaceInput(output2, input); + input->addTarget(consumer); + // Update predecessor/successor: transpose2 is removed, so consumer loses it as predecessor + consumer->removePredecessors(transpose2); + // If input has a source, it becomes a new predecessor of consumer + if (auto inputSource = input->getSource()) { + inputSource->addSuccessors(consumer); + consumer->addPredecessors(inputSource); + } + } + input->removeTarget(transpose1); + + // Clean up the graph connections for the removed operators + transpose1->removeSuccessors(transpose2); + transpose2->removePredecessors(transpose1); + + // Remove tensors + output->setSource(nullptr); + output->removeTarget(transpose2); + output2->setSource(nullptr); + for (auto &t : output2->getTargets()) + output2->removeTarget(t); + removeTensor(output); + removeTensor(output2); + + // Remove both ops + it = ops.erase(it); + for (auto it2 = ops.begin(); it2 != ops.end(); ++it2) { + if (*it2 == nextOp) { + ops.erase(it2); + break; + } + } + changed = true; + break; // Restart iteration after modification + } + if (changed) + continue; + + // Rule 2: Fuse transpose into matmul + for (auto it = ops.begin(); it != ops.end(); ) { + auto op = *it; + if (op->getOpType() != OpType::Transpose) { + ++it; + continue; + } + auto transpose = as(op); + auto output = transpose->getOutput(); + auto targets = output->getTargets(); + if (targets.size() != 1) { + ++it; + continue; + } + auto consumer = targets[0]; + if (consumer->getOpType() != OpType::MatMul) { + ++it; + continue; + } + const auto &perm = transpose->getPermute(); + int rank = perm.size(); + // Check if transpose only swaps the last two dimensions + bool swapsLastTwo = true; + for (int i = 0; i < rank - 2; ++i) { + if (perm[i] != i) { + swapsLastTwo = false; + break; + } + } + if (!swapsLastTwo || perm[rank - 2] != rank - 1 || perm[rank - 1] != rank - 2) { + ++it; + continue; + } + // Fuse into matmul + auto matmul = as(consumer); + auto input = transpose->getInputs(0); + if (output == matmul->getInputs(0)) { + matmul->setTransA(!matmul->getTransA()); + } else if (output == matmul->getInputs(1)) { + matmul->setTransB(!matmul->getTransB()); + } + // Replace output with input in matmul + matmul->replaceInput(output, input); + input->removeTarget(transpose); + input->addTarget(matmul); + + // Update predecessor/successor: transpose is removed, so matmul loses it as predecessor + matmul->removePredecessors(transpose); + transpose->removeSuccessors(matmul); + // If input has a source, it becomes a new predecessor of matmul + if (auto inputSource = input->getSource()) { + inputSource->addSuccessors(matmul); + matmul->addPredecessors(inputSource); + } + + // Remove transpose and output tensor + output->setSource(nullptr); + for (auto &t : output->getTargets()) + output->removeTarget(t); + removeTensor(output); + it = ops.erase(it); + changed = true; + break; + } + } + } + + Tensor GraphObj::getTensor(int fuid) const + { + for (auto tensor : tensors) + { + if (tensor->getFuid() == fuid) + { + return tensor; + } + } + return nullptr; + } + + void GraphObj::shape_infer() + { + for (auto &op : ops) + { + auto ans = op->inferShape(); + IT_ASSERT(ans.has_value()); + auto oldOutputs = op->getOutputs(); + IT_ASSERT(ans.value().size() == oldOutputs.size()); + // replace the old outputshape and size with new one + for (int i = 0; i < (int)ans.value().size(); ++i) + { + auto newShape = ans.value()[i]; + auto oldShape = oldOutputs[i]->getDims(); + auto fuid = oldOutputs[i]->getFuid(); + if (newShape != oldShape) + { + auto tensor = this->getTensor(fuid); + tensor->setShape(newShape); + } + } + } + } + + void GraphObj::dataMalloc() + { + // topological sorting first + IT_ASSERT(topo_sort() == true); + + // =================================== 作业 =================================== + // TODO：利用 allocator 给计算图分配内存 + // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存 + // =================================== 作业 =================================== + + // Track tensor liveness: for each tensor, find the index of its last consumer + std::unordered_map lastConsumerIdx; + std::unordered_map consumerCount; + for (auto &tensor : tensors) { + consumerCount[tensor->getGuid()] = tensor->getTargets().size(); + lastConsumerIdx[tensor->getGuid()] = -1; + } + + // Find the index of each operator in topological order + std::unordered_map opIdx; + for (int i = 0; i < (int)ops.size(); ++i) { + opIdx[ops[i]->getGuid()] = i; + } + + // For each tensor, find the last consumer index + for (auto &tensor : tensors) { + for (auto &target : tensor->getTargets()) { + int idx = opIdx[target->getGuid()]; + lastConsumerIdx[tensor->getGuid()] = std::max(lastConsumerIdx[tensor->getGuid()], idx); + } + // If tensor has no consumers (graph output), keep it alive + if (tensor->getTargets().empty()) { + if (auto src = tensor->getSource()) { + lastConsumerIdx[tensor->getGuid()] = opIdx[src->getGuid()]; + } + } + } + + // Allocate memory for each tensor, freeing when no longer needed + std::unordered_map tensorOffsets; + + // First, allocate graph input tensors (tensors without a source) + for (auto &tensor : tensors) { + if (!tensor->getSource() && tensor->getTargets().size() > 0) { + size_t offset = allocator.alloc(tensor->getBytes()); + tensorOffsets[tensor->getGuid()] = offset; + } + } + + for (int i = 0; i < (int)ops.size(); ++i) { + auto &op = ops[i]; + // Allocate output tensors + for (auto &output : op->getOutputs()) { + size_t offset = allocator.alloc(output->getBytes()); + tensorOffsets[output->getGuid()] = offset; + } + // Free input tensors that are no longer needed + for (auto &input : op->getInputs()) { + if (input) { + int remaining = --consumerCount[input->getGuid()]; + if (remaining == 0 && lastConsumerIdx[input->getGuid()] == i) { + allocator.free(tensorOffsets[input->getGuid()], input->getBytes()); + } + } + } + } + + // Get the actual memory pointer and assign Blobs to tensors + void *basePtr = allocator.getPtr(); + for (auto &tensor : tensors) { + if (tensorOffsets.count(tensor->getGuid())) { + auto blob = make_ref(runtime, (char *)basePtr + tensorOffsets[tensor->getGuid()]); + tensor->setDataBlob(blob); + } + } + + allocator.info(); + } + + Tensor GraphObj::addTensor(Shape dim, DataType dtype) + { + return tensors.emplace_back(make_ref(dim, dtype, runtime)); + } + + Tensor GraphObj::addTensor(const Tensor &tensor) + { + IT_ASSERT(tensor->getRuntime() == runtime, + std::string("Tensor runtime mismatch: cannot add a tenosr in ") + + tensor->getRuntime()->toString() + " to " + + runtime->toString()); + tensors.emplace_back(tensor); + return tensor; + } + + TensorVec GraphObj::addTensor(const TensorVec &tensors) + { + for (auto &t : tensors) + addTensor(t); + return tensors; + } + + // tensor's "source" and "target" must be in "ops". + // tensor has no "source" and no "target" must not exist. + // "inputs" or "outputs" of operators must be in "tensors" + // "predecessors" and "successors" of an operator of "ops" must be in "ops". + bool GraphObj::checkValid() const + { + for (auto tensor : tensors) + { + IT_ASSERT(!(tensor->getTargets().size() == 0 && + nullptr == tensor->getSource())); + for (auto op : tensor->getTargets()) + { + IT_ASSERT(std::find(ops.begin(), ops.end(), op) != ops.end()); + } + auto op = tensor->getSource(); + IT_ASSERT(!(op && std::find(ops.begin(), ops.end(), op) == ops.end())); + } + for (auto op : ops) + { + for (auto tensor : op->getInputs()) + { + IT_ASSERT(std::find(tensors.begin(), tensors.end(), tensor) != + tensors.end()); + } + for (auto tensor : op->getOutputs()) + { + IT_ASSERT(std::find(tensors.begin(), tensors.end(), tensor) != + tensors.end()); + } + for (auto pre : op->getPredecessors()) + { + IT_ASSERT(std::find(ops.begin(), ops.end(), pre) != ops.end()); + } + for (auto suc : op->getSuccessors()) + { + IT_ASSERT(std::find(ops.begin(), ops.end(), suc) != ops.end()); + } + } + std::set s; + // check whether two tensors with the same FUID exist + for (auto tensor : tensors) + { + int cnt = s.count(tensor->getFuid()); + IT_ASSERT(cnt == 0, std::to_string(tensor->getFuid())); + s.insert(tensor->getFuid()); + } + return true; + } + } // namespace infini \ No newline at end of file diff --git a/src/core/op_type.cc b/src/core/op_type.cc index b2a721a3..b2af61b2 100644 --- a/src/core/op_type.cc +++ b/src/core/op_type.cc @@ -1,32 +1,32 @@ -#include "core/op_type.h" - -namespace infini -{ - const char *OpType::toString() const - { -#define CASE(NAME) \ - case OpType::NAME: \ - return #NAME - - switch (type) - { - CASE(Unknown); - CASE(Add); - CASE(Sub); - CASE(Mul); - CASE(Div); - CASE(Cast); - CASE(Clip); - CASE(Relu); - CASE(Transpose); - CASE(Concat); - CASE(MatMul); - - default: - return "Unknown"; - } - -#undef CASE - } - -} // namespace infini +#include "core/op_type.h" + +namespace infini +{ + const char *OpType::toString() const + { +#define CASE(NAME) \ + case OpType::NAME: \ + return #NAME + + switch (type) + { + CASE(Unknown); + CASE(Add); + CASE(Sub); + CASE(Mul); + CASE(Div); + CASE(Cast); + CASE(Clip); + CASE(Relu); + CASE(Transpose); + CASE(Concat); + CASE(MatMul); + + default: + return "Unknown"; + } + +#undef CASE + } + +} // namespace infini diff --git a/src/core/operator.cc b/src/core/operator.cc index a70ca48c..d3622fb2 100644 --- a/src/core/operator.cc +++ b/src/core/operator.cc @@ -1,85 +1,85 @@ -#include "core/operator.h" -#include "core/graph.h" - -namespace infini -{ - - OperatorObj::OperatorObj(OpType opType, TensorVec inputs, TensorVec outputs) - : type(opType), inputs(inputs), outputs(outputs) {} - - void OperatorObj::removePredecessors(const Operator &op) - { - for (auto it = predecessors.begin(); it != predecessors.end();) - { - if (it->lock() == op) - it = predecessors.erase(it); - else - ++it; - } - } - - void OperatorObj::removeSuccessors(const Operator &op) - { - for (auto it = successors.begin(); it != successors.end();) - { - if (it->lock() == op) - it = successors.erase(it); - else - ++it; - } - } - - void OperatorObj::replaceInput(Tensor t1, Tensor t2) - { - for (auto itr = inputs.begin(); itr != inputs.end(); ++itr) - { - if (*itr == t1) - { - *itr = t2; - } - } - } - - bool OperatorObj::checkValid(GraphObj *graph) - { - auto optShapes = inferShape(); - if (!optShapes) // shape inference failed - return false; - - const vector &shapes = *optShapes; - if (shapes.size() != outputs.size()) - return false; - if (graph) - { // if graph != nullptr, outputs should be created - auto dataTypes = inferDataType(); - for (size_t i = 0; i < outputs.size(); i++) - { - IT_ASSERT(!outputs[i], "Find empty output while operator creation"); - outputs[i] = graph->addTensor(shapes[i], dataTypes[i]); - } - } - else - { // if outputs have been created, check their shapes - for (size_t i = 0; i < shapes.size(); ++i) - { - if (shapes[i] != outputs[i]->getDims()) - return false; - } - } - return true; - } - - optional> OperatorObj::inferShape() { return inferShape(inputs); } - - vector OperatorObj::inferDataType(const TensorVec &inputs) const - { - auto dataType = inputs[0]->getDType(); - return vector(numOutputs(), dataType); - } - - vector OperatorObj::inferDataType() const - { - return inferDataType(inputs); - } - -} // namespace infini +#include "core/operator.h" +#include "core/graph.h" + +namespace infini +{ + + OperatorObj::OperatorObj(OpType opType, TensorVec inputs, TensorVec outputs) + : type(opType), inputs(inputs), outputs(outputs) {} + + void OperatorObj::removePredecessors(const Operator &op) + { + for (auto it = predecessors.begin(); it != predecessors.end();) + { + if (it->lock() == op) + it = predecessors.erase(it); + else + ++it; + } + } + + void OperatorObj::removeSuccessors(const Operator &op) + { + for (auto it = successors.begin(); it != successors.end();) + { + if (it->lock() == op) + it = successors.erase(it); + else + ++it; + } + } + + void OperatorObj::replaceInput(Tensor t1, Tensor t2) + { + for (auto itr = inputs.begin(); itr != inputs.end(); ++itr) + { + if (*itr == t1) + { + *itr = t2; + } + } + } + + bool OperatorObj::checkValid(GraphObj *graph) + { + auto optShapes = inferShape(); + if (!optShapes) // shape inference failed + return false; + + const vector &shapes = *optShapes; + if (shapes.size() != outputs.size()) + return false; + if (graph) + { // if graph != nullptr, outputs should be created + auto dataTypes = inferDataType(); + for (size_t i = 0; i < outputs.size(); i++) + { + IT_ASSERT(!outputs[i], "Find empty output while operator creation"); + outputs[i] = graph->addTensor(shapes[i], dataTypes[i]); + } + } + else + { // if outputs have been created, check their shapes + for (size_t i = 0; i < shapes.size(); ++i) + { + if (shapes[i] != outputs[i]->getDims()) + return false; + } + } + return true; + } + + optional> OperatorObj::inferShape() { return inferShape(inputs); } + + vector OperatorObj::inferDataType(const TensorVec &inputs) const + { + auto dataType = inputs[0]->getDType(); + return vector(numOutputs(), dataType); + } + + vector OperatorObj::inferDataType() const + { + return inferDataType(inputs); + } + +} // namespace infini diff --git a/src/core/runtime.cc b/src/core/runtime.cc index bd88d904..09692612 100644 --- a/src/core/runtime.cc +++ b/src/core/runtime.cc @@ -1,36 +1,36 @@ -#include "core/runtime.h" -#include "core/blob.h" -#include "core/kernel.h" -#include "core/graph.h" -#include "core/kernel.h" -#include -#include -#include -namespace infini -{ - void NativeCpuRuntimeObj::run(const Graph &graph) const - { - const auto &kernelRegistry = KernelRegistry::getInstance(); - - for (auto &op : graph->getOperators()) - { - auto kernelAttrs = KernelAttrs{device, op->getOpType().underlying()}; - Kernel *kernel = kernelRegistry.getKernel(kernelAttrs); - kernel->compute(op, this); - } - } - - string NativeCpuRuntimeObj::toString() const { return "CPU Runtime"; } - - void NativeCpuRuntimeObj::dealloc(void *ptr) - { - return free(ptr); - } - - void *NativeCpuRuntimeObj::alloc(size_t size) - { - return calloc((size + sizeof(uint64_t) - 1) / sizeof(uint64_t), - sizeof(uint64_t)); - } - -} // namespace infini +#include "core/runtime.h" +#include "core/blob.h" +#include "core/kernel.h" +#include "core/graph.h" +#include "core/kernel.h" +#include +#include +#include +namespace infini +{ + void NativeCpuRuntimeObj::run(const Graph &graph) const + { + const auto &kernelRegistry = KernelRegistry::getInstance(); + + for (auto &op : graph->getOperators()) + { + auto kernelAttrs = KernelAttrs{device, op->getOpType().underlying()}; + Kernel *kernel = kernelRegistry.getKernel(kernelAttrs); + kernel->compute(op, this); + } + } + + string NativeCpuRuntimeObj::toString() const { return "CPU Runtime"; } + + void NativeCpuRuntimeObj::dealloc(void *ptr) + { + return free(ptr); + } + + void *NativeCpuRuntimeObj::alloc(size_t size) + { + return calloc((size + sizeof(uint64_t) - 1) / sizeof(uint64_t), + sizeof(uint64_t)); + } + +} // namespace infini diff --git a/src/core/tensor.cc b/src/core/tensor.cc index db54a2d6..6c11869a 100644 --- a/src/core/tensor.cc +++ b/src/core/tensor.cc @@ -1,116 +1,116 @@ -#include "core/tensor.h" -#include "core/blob.h" -#include "core/operator.h" -#include "core/runtime.h" -#include -#include - -namespace infini { - - TensorObj::TensorObj(Shape shape_, DataType dtype, Runtime runtime) - : dim(shape_.size()), dtype(dtype), runtime(runtime), shape(std::move(shape_)), - _size(std::accumulate(shape.begin(), shape.end(), 1, std::multiplies{})) {} - - string TensorObj::toString() const - { - // Convert data pointer to string - std::stringstream ss; - if (data != nullptr) - ss << data->getPtr(); - else - ss << "nullptr data"; - string ret = "Tensor " + std::to_string(guid) + ", Fuid " + - std::to_string(fuid) + ", shape " + vecToString(shape) + - ", dtype " + dtype.toString() + ", " + runtime->toString() + - ", " + ss.str() + "\n"; - vector targetGuids; - for (const auto &op : targets) - targetGuids.emplace_back(op.lock()->getGuid()); - if (auto o = source.lock()) - ret += ", source " + std::to_string(o->getGuid()); - else - ret += ", source None"; - ret += ", targets " + vecToString(targetGuids); - return ret; - } - -void TensorObj::setShape(Shape shape_) { - shape = shape_; - size_t size = std::accumulate(shape.begin(), shape.end(), 1, - [](auto acc, auto x) { return acc * x; }); - _size = size; -} - -void TensorObj::printData() const { - IT_ASSERT(data != nullptr); - if (!runtime->isCpu()) - IT_TODO_HALT(); - -#define TRY_PRINT(N) \ - if (dtype == DataType(N)) \ - std::cout << dataToString::t>() << std::endl; - - TRY_PRINT(0) // fmt: new line - else TRY_PRINT(1) // - else TRY_PRINT(2) // - else TRY_PRINT(3) // - else TRY_PRINT(4) // - else TRY_PRINT(5) // - else TRY_PRINT(6) // - else TRY_PRINT(7) // - else TRY_PRINT(8) // - else TRY_PRINT(9) // - else TRY_PRINT(10) // - else TRY_PRINT(11) // - else TRY_PRINT(12) // - else TRY_PRINT(13) // - else TRY_PRINT(16) // - else IT_TODO_HALT(); - -#undef TRY_PRINT -} - -bool TensorObj::equalData(const Tensor &rhs, double relativeError) const { - IT_ASSERT(data != nullptr); - IT_ASSERT(rhs->data != nullptr); - IT_ASSERT(getDType() == rhs->getDType()); - IT_ASSERT(runtime->isCpu()); - IT_ASSERT(rhs->getRuntime()->isCpu()); - if (size() != rhs->size()) - return false; - -#define TEST_EQUAL(N) \ - if (dtype == DataType(N)) \ - return equalDataImpl(getRawDataPtr::t *>(), \ - rhs->getRawDataPtr::t *>(), size(), \ - relativeError); - - TEST_EQUAL(0) // fmt: new line - else TEST_EQUAL(1) // - else TEST_EQUAL(2) // - else TEST_EQUAL(3) // - else TEST_EQUAL(4) // - else TEST_EQUAL(5) // - else TEST_EQUAL(6) // - else TEST_EQUAL(7) // - else TEST_EQUAL(8) // - else TEST_EQUAL(9) // - else TEST_EQUAL(10) // - else TEST_EQUAL(11) // - else TEST_EQUAL(12) // - else TEST_EQUAL(13) // - else TEST_EQUAL(16) // - else IT_TODO_HALT(); - -#undef TEST_EQUAL -} - -void TensorObj::setData( - const std::function &generator) const { - IT_ASSERT(data != nullptr); - generator(getRawDataPtr(), size(), dtype); -} - -void TensorObj::setDataBlob(const Blob &blob) { this->data = blob; } - -}; // namespace infini +#include "core/tensor.h" +#include "core/blob.h" +#include "core/operator.h" +#include "core/runtime.h" +#include +#include + +namespace infini { + + TensorObj::TensorObj(Shape shape_, DataType dtype, Runtime runtime) + : dim(shape_.size()), dtype(dtype), runtime(runtime), shape(std::move(shape_)), + _size(std::accumulate(shape.begin(), shape.end(), 1, std::multiplies{})) {} + + string TensorObj::toString() const + { + // Convert data pointer to string + std::stringstream ss; + if (data != nullptr) + ss << data->getPtr(); + else + ss << "nullptr data"; + string ret = "Tensor " + std::to_string(guid) + ", Fuid " + + std::to_string(fuid) + ", shape " + vecToString(shape) + + ", dtype " + dtype.toString() + ", " + runtime->toString() + + ", " + ss.str() + "\n"; + vector targetGuids; + for (const auto &op : targets) + targetGuids.emplace_back(op.lock()->getGuid()); + if (auto o = source.lock()) + ret += ", source " + std::to_string(o->getGuid()); + else + ret += ", source None"; + ret += ", targets " + vecToString(targetGuids); + return ret; + } + +void TensorObj::setShape(Shape shape_) { + shape = shape_; + size_t size = std::accumulate(shape.begin(), shape.end(), 1, + [](auto acc, auto x) { return acc * x; }); + _size = size; +} + +void TensorObj::printData() const { + IT_ASSERT(data != nullptr); + if (!runtime->isCpu()) + IT_TODO_HALT(); + +#define TRY_PRINT(N) \ + if (dtype == DataType(N)) \ + std::cout << dataToString::t>() << std::endl; + + TRY_PRINT(0) // fmt: new line + else TRY_PRINT(1) // + else TRY_PRINT(2) // + else TRY_PRINT(3) // + else TRY_PRINT(4) // + else TRY_PRINT(5) // + else TRY_PRINT(6) // + else TRY_PRINT(7) // + else TRY_PRINT(8) // + else TRY_PRINT(9) // + else TRY_PRINT(10) // + else TRY_PRINT(11) // + else TRY_PRINT(12) // + else TRY_PRINT(13) // + else TRY_PRINT(16) // + else IT_TODO_HALT(); + +#undef TRY_PRINT +} + +bool TensorObj::equalData(const Tensor &rhs, double relativeError) const { + IT_ASSERT(data != nullptr); + IT_ASSERT(rhs->data != nullptr); + IT_ASSERT(getDType() == rhs->getDType()); + IT_ASSERT(runtime->isCpu()); + IT_ASSERT(rhs->getRuntime()->isCpu()); + if (size() != rhs->size()) + return false; + +#define TEST_EQUAL(N) \ + if (dtype == DataType(N)) \ + return equalDataImpl(getRawDataPtr::t *>(), \ + rhs->getRawDataPtr::t *>(), size(), \ + relativeError); + + TEST_EQUAL(0) // fmt: new line + else TEST_EQUAL(1) // + else TEST_EQUAL(2) // + else TEST_EQUAL(3) // + else TEST_EQUAL(4) // + else TEST_EQUAL(5) // + else TEST_EQUAL(6) // + else TEST_EQUAL(7) // + else TEST_EQUAL(8) // + else TEST_EQUAL(9) // + else TEST_EQUAL(10) // + else TEST_EQUAL(11) // + else TEST_EQUAL(12) // + else TEST_EQUAL(13) // + else TEST_EQUAL(16) // + else IT_TODO_HALT(); + +#undef TEST_EQUAL +} + +void TensorObj::setData( + const std::function &generator) const { + IT_ASSERT(data != nullptr); + generator(getRawDataPtr(), size(), dtype); +} + +void TensorObj::setDataBlob(const Blob &blob) { this->data = blob; } + +}; // namespace infini diff --git a/src/kernels/cpu/concat.cc b/src/kernels/cpu/concat.cc index 6e061c75..2c511281 100644 --- a/src/kernels/cpu/concat.cc +++ b/src/kernels/cpu/concat.cc @@ -1,64 +1,64 @@ -#include "operators/concat.h" -#include "core/kernel.h" - -namespace infini { - -class NaiveConcat : public CpuKernelWithoutConfig { - template - void doCompute(const Operator &_op, const RuntimeObj *context) const { - auto op = as(_op); - auto inputs = op->getInputs(), outputs = op->getOutputs(); - auto dim = op->getDim(); - auto output = outputs[0]; - std::vector iDims; - for (auto input : inputs) - iDims.emplace_back(input->getDims()); - const auto &outDim = output->getDims(); - size_t blockOffsetInner = 1; - for (size_t i = outDim.size() - 1; i > (size_t)dim; --i) - blockOffsetInner *= outDim[i]; - size_t blockOffset = outDim[dim] * blockOffsetInner; - for (size_t i = 0; i < inputs.size(); ++i) { - auto input = inputs[i]; - auto dimOffset = 0; - auto iDim = iDims[i]; - for (size_t j = 0; j < i; ++j) - dimOffset += iDims[j][dim]; - size_t localBlockOffset = 1; - for (size_t i = iDim.size() - 1; - i >= (size_t)dim && i != (size_t)-1; --i) - localBlockOffset *= iDim[i]; - auto innerOffset = blockOffsetInner * dimOffset; - auto inSize = input->size(); - auto inPtr = input->getRawDataPtr(), - outPtr = output->getRawDataPtr(); -#pragma omp parallel for - for (size_t iOffset = 0; iOffset < inSize; ++iOffset) { - auto oOffset = iOffset % localBlockOffset + innerOffset + - iOffset / localBlockOffset * blockOffset; - outPtr[oOffset] = inPtr[iOffset]; - } - } - } - - void compute(const Operator &_op, - const RuntimeObj *context) const override { -#define CASE(N) \ - case N: \ - doCompute::t>(_op, context) - - int dataTypeIdx = _op->getDType().getIndex(); - switch (dataTypeIdx) { - CASE(1); // DataType::Float32 - break; - CASE(12); // DataType::UInt32 - break; - default: - IT_TODO_HALT(); - } - } -}; - -REGISTER_KERNEL(Device::CPU, OpType::Concat, NaiveConcat, "ConcatNaive_CPU"); - -} // namespace infini +#include "operators/concat.h" +#include "core/kernel.h" + +namespace infini { + +class NaiveConcat : public CpuKernelWithoutConfig { + template + void doCompute(const Operator &_op, const RuntimeObj *context) const { + auto op = as(_op); + auto inputs = op->getInputs(), outputs = op->getOutputs(); + auto dim = op->getDim(); + auto output = outputs[0]; + std::vector iDims; + for (auto input : inputs) + iDims.emplace_back(input->getDims()); + const auto &outDim = output->getDims(); + size_t blockOffsetInner = 1; + for (size_t i = outDim.size() - 1; i > (size_t)dim; --i) + blockOffsetInner *= outDim[i]; + size_t blockOffset = outDim[dim] * blockOffsetInner; + for (size_t i = 0; i < inputs.size(); ++i) { + auto input = inputs[i]; + auto dimOffset = 0; + auto iDim = iDims[i]; + for (size_t j = 0; j < i; ++j) + dimOffset += iDims[j][dim]; + size_t localBlockOffset = 1; + for (size_t i = iDim.size() - 1; + i >= (size_t)dim && i != (size_t)-1; --i) + localBlockOffset *= iDim[i]; + auto innerOffset = blockOffsetInner * dimOffset; + auto inSize = input->size(); + auto inPtr = input->getRawDataPtr(), + outPtr = output->getRawDataPtr(); +#pragma omp parallel for + for (size_t iOffset = 0; iOffset < inSize; ++iOffset) { + auto oOffset = iOffset % localBlockOffset + innerOffset + + iOffset / localBlockOffset * blockOffset; + outPtr[oOffset] = inPtr[iOffset]; + } + } + } + + void compute(const Operator &_op, + const RuntimeObj *context) const override { +#define CASE(N) \ + case N: \ + doCompute::t>(_op, context) + + int dataTypeIdx = _op->getDType().getIndex(); + switch (dataTypeIdx) { + CASE(1); // DataType::Float32 + break; + CASE(12); // DataType::UInt32 + break; + default: + IT_TODO_HALT(); + } + } +}; + +REGISTER_KERNEL(Device::CPU, OpType::Concat, NaiveConcat, "ConcatNaive_CPU"); + +} // namespace infini diff --git a/src/kernels/cpu/element_wise.cc b/src/kernels/cpu/element_wise.cc index af03c7a3..ba3361cd 100644 --- a/src/kernels/cpu/element_wise.cc +++ b/src/kernels/cpu/element_wise.cc @@ -1,119 +1,119 @@ -#include "operators/element_wise.h" -#include "core/kernel.h" -#include "utils/operator_utils.h" - -namespace infini -{ - class NativeElementWise : public CpuKernelWithoutConfig - { - template - static T addCompute(T val0, T val1) - { - return val0 + val1; - } - - template - static T subCompute(T val0, T val1) - { - return val0 - val1; - } - - template - static T mulCompute(T val0, T val1) - { - return val0 * val1; - } - - template - static T divCompute(T val0, T val1) - { - return (T)(val0 / val1); - } - - template - void doCompute(const Operator &_op, const RuntimeObj *context) const - { - auto op = as(_op); - T *inptr0 = op->getInputs(0)->getRawDataPtr(); - T *inptr1 = op->getInputs(1)->getRawDataPtr(); - T *outptr = op->getOutput()->getRawDataPtr(); - - auto shapeA = op->getInputs(0)->getDims(); - auto shapeB = op->getInputs(1)->getDims(); - auto shapeC = op->getOutput()->getDims(); - auto rank = op->getOutput()->getRank(); - Shape a(rank, 1); - Shape b(rank, 1); - std::copy(shapeA.begin(), shapeA.end(), - a.begin() + (rank - shapeA.size())); - std::copy(shapeB.begin(), shapeB.end(), - b.begin() + (rank - shapeB.size())); - auto getStride = [&](const Shape &shape) - { - int p = 1; - Shape stride(rank); - for (auto i = rank; i > 0; --i) - { - stride[i - 1] = p; - p = p * shape[i - 1]; - } - return stride; - }; - Shape strideA = getStride(a); - Shape strideB = getStride(b); - - auto n = op->getOutput()->size(); - T (*_doCompute) - (T val0, T val1); - switch (op->getOpType().underlying()) - { - case OpType::Add: - _doCompute = addCompute; - break; - case OpType::Sub: - _doCompute = subCompute; - break; - case OpType::Mul: - _doCompute = mulCompute; - break; - case OpType::Div: - _doCompute = divCompute; - break; - default: - IT_TODO_HALT(); - } - - for (size_t i = 0; i < n; ++i) - { - auto shapeIndexC = locate_index(i, shapeC); - auto indexA = delocate_index(shapeIndexC, a, strideA); - auto indexB = delocate_index(shapeIndexC, b, strideB); - outptr[i] = _doCompute(inptr0[indexA], inptr1[indexB]); - } - } - - void compute(const Operator &_op, - const RuntimeObj *context) const override - { -#define CASE(N) \ - case N: \ - doCompute::t>(_op, context) - - int dataTypeIdx = _op->getDType().getIndex(); - switch (dataTypeIdx) - { - CASE(1); // DataType::Float32 - break; - CASE(12); // DataType::UInt32 - break; - default: - IT_TODO_HALT(); - } - } - }; - - REGISTER_KERNEL(Device::CPU, OpType::Add, NativeElementWise, "addNaive_CPU"); - REGISTER_KERNEL(Device::CPU, OpType::Sub, NativeElementWise, "subNaive_CPU"); - REGISTER_KERNEL(Device::CPU, OpType::Mul, NativeElementWise, "mulNaive_CPU"); - REGISTER_KERNEL(Device::CPU, OpType::Div, NativeElementWise, "divNaive_CPU"); -}; // namespace infini +#include "operators/element_wise.h" +#include "core/kernel.h" +#include "utils/operator_utils.h" + +namespace infini +{ + class NativeElementWise : public CpuKernelWithoutConfig + { + template + static T addCompute(T val0, T val1) + { + return val0 + val1; + } + + template + static T subCompute(T val0, T val1) + { + return val0 - val1; + } + + template + static T mulCompute(T val0, T val1) + { + return val0 * val1; + } + + template + static T divCompute(T val0, T val1) + { + return (T)(val0 / val1); + } + + template + void doCompute(const Operator &_op, const RuntimeObj *context) const + { + auto op = as(_op); + T *inptr0 = op->getInputs(0)->getRawDataPtr(); + T *inptr1 = op->getInputs(1)->getRawDataPtr(); + T *outptr = op->getOutput()->getRawDataPtr(); + + auto shapeA = op->getInputs(0)->getDims(); + auto shapeB = op->getInputs(1)->getDims(); + auto shapeC = op->getOutput()->getDims(); + auto rank = op->getOutput()->getRank(); + Shape a(rank, 1); + Shape b(rank, 1); + std::copy(shapeA.begin(), shapeA.end(), + a.begin() + (rank - shapeA.size())); + std::copy(shapeB.begin(), shapeB.end(), + b.begin() + (rank - shapeB.size())); + auto getStride = [&](const Shape &shape) + { + int p = 1; + Shape stride(rank); + for (auto i = rank; i > 0; --i) + { + stride[i - 1] = p; + p = p * shape[i - 1]; + } + return stride; + }; + Shape strideA = getStride(a); + Shape strideB = getStride(b); + + auto n = op->getOutput()->size(); + T (*_doCompute) + (T val0, T val1); + switch (op->getOpType().underlying()) + { + case OpType::Add: + _doCompute = addCompute; + break; + case OpType::Sub: + _doCompute = subCompute; + break; + case OpType::Mul: + _doCompute = mulCompute; + break; + case OpType::Div: + _doCompute = divCompute; + break; + default: + IT_TODO_HALT(); + } + + for (size_t i = 0; i < n; ++i) + { + auto shapeIndexC = locate_index(i, shapeC); + auto indexA = delocate_index(shapeIndexC, a, strideA); + auto indexB = delocate_index(shapeIndexC, b, strideB); + outptr[i] = _doCompute(inptr0[indexA], inptr1[indexB]); + } + } + + void compute(const Operator &_op, + const RuntimeObj *context) const override + { +#define CASE(N) \ + case N: \ + doCompute::t>(_op, context) + + int dataTypeIdx = _op->getDType().getIndex(); + switch (dataTypeIdx) + { + CASE(1); // DataType::Float32 + break; + CASE(12); // DataType::UInt32 + break; + default: + IT_TODO_HALT(); + } + } + }; + + REGISTER_KERNEL(Device::CPU, OpType::Add, NativeElementWise, "addNaive_CPU"); + REGISTER_KERNEL(Device::CPU, OpType::Sub, NativeElementWise, "subNaive_CPU"); + REGISTER_KERNEL(Device::CPU, OpType::Mul, NativeElementWise, "mulNaive_CPU"); + REGISTER_KERNEL(Device::CPU, OpType::Div, NativeElementWise, "divNaive_CPU"); +}; // namespace infini diff --git a/src/kernels/cpu/transpose.cc b/src/kernels/cpu/transpose.cc index 46292d45..b43dd6ec 100644 --- a/src/kernels/cpu/transpose.cc +++ b/src/kernels/cpu/transpose.cc @@ -1,60 +1,60 @@ -#include "operators/transpose.h" -#include "core/kernel.h" - -namespace infini { - -inline Shape idx2Pos(const Shape &shape, size_t idx) { - Shape pos = Shape(shape.size(), 0); - auto rest = idx, curDimId = shape.size() - 1; - while (rest > 0) { - pos[curDimId] = rest % shape[curDimId]; - rest /= shape[curDimId]; - curDimId--; - } - return pos; -} - -class NaiveTranspose : public CpuKernelWithoutConfig { - template - void doCompute(const Operator &_op, const RuntimeObj *context) const { - auto op = as(_op); - auto inputs = op->getInputs(), outputs = op->getOutputs(); - const auto &inDim = inputs[0]->getDims(); - const auto &perm = op->getPermute(); - - size_t inSize = inputs[0]->size(); - auto inPtr = inputs[0]->getRawDataPtr(), - outPtr = outputs[0]->getRawDataPtr(); - // #pragma omp parallel for - for (size_t inIdx = 0; inIdx < inSize; ++inIdx) { - auto posInput = idx2Pos(inDim, inIdx); - int outIdx = 0; - for (size_t j = 0, jEnd = perm.size(); j < jEnd; ++j) { - outIdx = outIdx * inDim[perm[j]] + posInput[perm[j]]; - } - outPtr[outIdx] = inPtr[inIdx]; - } - } - - void compute(const Operator &_op, - const RuntimeObj *context) const override { -#define CASE(N) \ - case N: \ - doCompute::t>(_op, context) - - int dataTypeIdx = _op->getDType().getIndex(); - switch (dataTypeIdx) { - CASE(1); // DataType::Float32 - break; - CASE(12); // DataType::UInt32 - break; - default: - IT_TODO_HALT(); - } - } -}; - -REGISTER_KERNEL(Device::CPU, OpType::Transpose, NaiveTranspose, - "TransposeNaive_CPU"); - -} // namespace infini +#include "operators/transpose.h" +#include "core/kernel.h" + +namespace infini { + +inline Shape idx2Pos(const Shape &shape, size_t idx) { + Shape pos = Shape(shape.size(), 0); + auto rest = idx, curDimId = shape.size() - 1; + while (rest > 0) { + pos[curDimId] = rest % shape[curDimId]; + rest /= shape[curDimId]; + curDimId--; + } + return pos; +} + +class NaiveTranspose : public CpuKernelWithoutConfig { + template + void doCompute(const Operator &_op, const RuntimeObj *context) const { + auto op = as(_op); + auto inputs = op->getInputs(), outputs = op->getOutputs(); + const auto &inDim = inputs[0]->getDims(); + const auto &perm = op->getPermute(); + + size_t inSize = inputs[0]->size(); + auto inPtr = inputs[0]->getRawDataPtr(), + outPtr = outputs[0]->getRawDataPtr(); + // #pragma omp parallel for + for (size_t inIdx = 0; inIdx < inSize; ++inIdx) { + auto posInput = idx2Pos(inDim, inIdx); + int outIdx = 0; + for (size_t j = 0, jEnd = perm.size(); j < jEnd; ++j) { + outIdx = outIdx * inDim[perm[j]] + posInput[perm[j]]; + } + outPtr[outIdx] = inPtr[inIdx]; + } + } + + void compute(const Operator &_op, + const RuntimeObj *context) const override { +#define CASE(N) \ + case N: \ + doCompute::t>(_op, context) + + int dataTypeIdx = _op->getDType().getIndex(); + switch (dataTypeIdx) { + CASE(1); // DataType::Float32 + break; + CASE(12); // DataType::UInt32 + break; + default: + IT_TODO_HALT(); + } + } +}; + +REGISTER_KERNEL(Device::CPU, OpType::Transpose, NaiveTranspose, + "TransposeNaive_CPU"); + +} // namespace infini diff --git a/src/kernels/cpu/unary.cc b/src/kernels/cpu/unary.cc index 29f88548..9e631c5f 100644 --- a/src/kernels/cpu/unary.cc +++ b/src/kernels/cpu/unary.cc @@ -1,105 +1,105 @@ -#include "operators/unary.h" -#include "core/kernel.h" - -namespace infini -{ - class NativeUnary : public CpuKernelWithoutConfig - { - template - static T reluCompute(T val) - { - return std::max(T(0), val); - } - - template - void doCompute(const Operator &_op, const RuntimeObj *context) const - { - auto op = as(_op); - T *inptr = op->getInputs(0)->getRawDataPtr(); - T *outptr = op->getOutput()->getRawDataPtr(); - - auto outDim = op->getOutput()->getDims(); - auto n = op->getOutput()->size(); - - T (*_doCompute) - (T val); - switch (op->getOpType().underlying()) - { - case OpType::Relu: - _doCompute = reluCompute; - break; - default: - IT_TODO_HALT(); - } - - for (size_t offset = 0; offset < n; offset++) - { - outptr[offset] = _doCompute(inptr[offset]); - } - } - - void compute(const Operator &_op, - const RuntimeObj *context) const override - { -#define CASE(N) \ - case N: \ - doCompute::t>(_op, context) - - int dataTypeIdx = _op->getDType().getIndex(); - switch (dataTypeIdx) - { - CASE(1); // DataType::Float32 - break; - CASE(12); // DataType::UInt32 - break; - default: - IT_TODO_HALT(); - } - } - }; - - class Clip : public CpuKernelWithoutConfig - { - template - void doCompute(const Operator &_op, const RuntimeObj *context) const - { - auto op = as(_op); - T *inptr = op->getInputs(0)->getRawDataPtr(); - T *outptr = op->getOutput()->getRawDataPtr(); - auto minValue = op->getMin(); - auto maxValue = op->getMax(); - - auto n = op->getOutput()->size(); - for (size_t offset = 0; offset < n; offset++) - { - auto val = *inptr++; - *outptr++ = (minValue && val < *minValue) ? *minValue - : (maxValue && val > *maxValue) ? *maxValue - : val; - } - } - - void compute(const Operator &_op, - const RuntimeObj *context) const override - { -#define CASE(N) \ - case N: \ - doCompute::t>(_op, context) - - int dataTypeIdx = _op->getDType().getIndex(); - switch (dataTypeIdx) - { - CASE(1); // DataType::Float32 - break; - CASE(12); // DataType::UInt32 - break; - default: - IT_TODO_HALT(); - } - } - }; - - REGISTER_KERNEL(Device::CPU, OpType::Relu, NativeUnary, "reluNaive_CPU"); - REGISTER_KERNEL(Device::CPU, OpType::Clip, Clip, "Clip_CPU"); - -}; // namespace infini +#include "operators/unary.h" +#include "core/kernel.h" + +namespace infini +{ + class NativeUnary : public CpuKernelWithoutConfig + { + template + static T reluCompute(T val) + { + return std::max(T(0), val); + } + + template + void doCompute(const Operator &_op, const RuntimeObj *context) const + { + auto op = as(_op); + T *inptr = op->getInputs(0)->getRawDataPtr(); + T *outptr = op->getOutput()->getRawDataPtr(); + + auto outDim = op->getOutput()->getDims(); + auto n = op->getOutput()->size(); + + T (*_doCompute) + (T val); + switch (op->getOpType().underlying()) + { + case OpType::Relu: + _doCompute = reluCompute; + break; + default: + IT_TODO_HALT(); + } + + for (size_t offset = 0; offset < n; offset++) + { + outptr[offset] = _doCompute(inptr[offset]); + } + } + + void compute(const Operator &_op, + const RuntimeObj *context) const override + { +#define CASE(N) \ + case N: \ + doCompute::t>(_op, context) + + int dataTypeIdx = _op->getDType().getIndex(); + switch (dataTypeIdx) + { + CASE(1); // DataType::Float32 + break; + CASE(12); // DataType::UInt32 + break; + default: + IT_TODO_HALT(); + } + } + }; + + class Clip : public CpuKernelWithoutConfig + { + template + void doCompute(const Operator &_op, const RuntimeObj *context) const + { + auto op = as(_op); + T *inptr = op->getInputs(0)->getRawDataPtr(); + T *outptr = op->getOutput()->getRawDataPtr(); + auto minValue = op->getMin(); + auto maxValue = op->getMax(); + + auto n = op->getOutput()->size(); + for (size_t offset = 0; offset < n; offset++) + { + auto val = *inptr++; + *outptr++ = (minValue && val < *minValue) ? *minValue + : (maxValue && val > *maxValue) ? *maxValue + : val; + } + } + + void compute(const Operator &_op, + const RuntimeObj *context) const override + { +#define CASE(N) \ + case N: \ + doCompute::t>(_op, context) + + int dataTypeIdx = _op->getDType().getIndex(); + switch (dataTypeIdx) + { + CASE(1); // DataType::Float32 + break; + CASE(12); // DataType::UInt32 + break; + default: + IT_TODO_HALT(); + } + } + }; + + REGISTER_KERNEL(Device::CPU, OpType::Relu, NativeUnary, "reluNaive_CPU"); + REGISTER_KERNEL(Device::CPU, OpType::Clip, Clip, "Clip_CPU"); + +}; // namespace infini diff --git a/src/operators/concat.cc b/src/operators/concat.cc index d1963308..5c5712b0 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -1,38 +1,40 @@ -#include "operators/concat.h" -#include "utils/operator_utils.h" - -namespace infini { -ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim) - : OperatorObj(OpType::Concat, inputs, {output}) { - int rank = inputs[0]->getRank(); - dim = get_real_axis(_dim, rank); - IT_ASSERT(checkValid(graph)); -} - -optional> ConcatObj::inferShape(const TensorVec &inputs) { - 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 - // =================================== 作业 =================================== - - return {{dims}}; -} - -std::string ConcatObj::toString() const { - std::ostringstream os; - os << "Concat[" << getGuid() << "]"; - os << "("; - for (auto input : inputs) - os << vecToString(input->getDims()) << ","; - os << "dim=" << dim << ","; - os << "input="; - for (auto input : inputs) - os << input->getGuid() << ","; - os << "output=" << outputs[0]->getGuid() << ")"; - return os.str(); -} - -} // namespace infini +#include "operators/concat.h" +#include "utils/operator_utils.h" + +namespace infini { +ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim) + : OperatorObj(OpType::Concat, inputs, {output}) { + int rank = inputs[0]->getRank(); + dim = get_real_axis(_dim, rank); + IT_ASSERT(checkValid(graph)); +} + +optional> ConcatObj::inferShape(const TensorVec &inputs) { + Shape dims = inputs[0]->getDims(); + + // =================================== 作业 =================================== + // TODO：修改 dims，返回正确的 concat 后的 shape + // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13 + // =================================== 作业 =================================== + for (size_t i = 1; i < inputs.size(); ++i) { + dims[dim] += inputs[i]->getDims()[dim]; + } + + return {{dims}}; +} + +std::string ConcatObj::toString() const { + std::ostringstream os; + os << "Concat[" << getGuid() << "]"; + os << "("; + for (auto input : inputs) + os << vecToString(input->getDims()) << ","; + os << "dim=" << dim << ","; + os << "input="; + for (auto input : inputs) + os << input->getGuid() << ","; + os << "output=" << outputs[0]->getGuid() << ")"; + return os.str(); +} + +} // namespace infini diff --git a/src/operators/element_wise.cc b/src/operators/element_wise.cc index c1b4ef1f..f9adb8f6 100644 --- a/src/operators/element_wise.cc +++ b/src/operators/element_wise.cc @@ -1,33 +1,33 @@ -#include "operators/element_wise.h" -#include "utils/operator_utils.h" - -namespace infini -{ - ElementWiseObj::ElementWiseObj(OpType type, GraphObj *graph, Tensor input0, - Tensor input1, Tensor output) - : OperatorObj(type, {input0, input1}, {output}) - { - IT_ASSERT(checkValid(graph)); - } - - optional> ElementWiseObj::inferShape(const TensorVec &inputs) - { - const auto A = inputs[0], B = inputs[1]; - auto res = infer_broadcast(A->getDims(), B->getDims()); - return {{res}}; - } - - std::string ElementWiseObj::toString() const - { - std::ostringstream os; - os << type.toString() << "[" << getGuid() << "]"; - os << "("; - os << vecToString(inputs[0]->getDims()) << ","; - os << vecToString(inputs[1]->getDims()) << ","; - os << "input0=" << inputs[0]->getGuid() << ","; - os << "input1=" << inputs[1]->getGuid() << ","; - os << "output=" << outputs[0]->getGuid() << ")"; - return os.str(); - } - -}; // namespace infini +#include "operators/element_wise.h" +#include "utils/operator_utils.h" + +namespace infini +{ + ElementWiseObj::ElementWiseObj(OpType type, GraphObj *graph, Tensor input0, + Tensor input1, Tensor output) + : OperatorObj(type, {input0, input1}, {output}) + { + IT_ASSERT(checkValid(graph)); + } + + optional> ElementWiseObj::inferShape(const TensorVec &inputs) + { + const auto A = inputs[0], B = inputs[1]; + auto res = infer_broadcast(A->getDims(), B->getDims()); + return {{res}}; + } + + std::string ElementWiseObj::toString() const + { + std::ostringstream os; + os << type.toString() << "[" << getGuid() << "]"; + os << "("; + os << vecToString(inputs[0]->getDims()) << ","; + os << vecToString(inputs[1]->getDims()) << ","; + os << "input0=" << inputs[0]->getGuid() << ","; + os << "input1=" << inputs[1]->getGuid() << ","; + os << "output=" << outputs[0]->getGuid() << ")"; + return os.str(); + } + +}; // namespace infini diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..b1cb9c33 100644 --- a/src/operators/matmul.cc +++ b/src/operators/matmul.cc @@ -1,33 +1,64 @@ -#include "operators/matmul.h" - -namespace infini -{ - - MatmulObj::MatmulObj(GraphObj *graph, Tensor A, Tensor B, Tensor C, bool transA, - bool transB) - : OperatorObj(OpType::MatMul, TensorVec{A, B}, {C}), - transA(transA), transB(transB) - { - IT_ASSERT(checkValid(graph)); - } - - string MatmulObj::toString() const - { - std::ostringstream os; - os << "Matmul([" << (transA ? "A^T" : "A") << "," << (transB ? "B^T" : "B]") - << ",A=" << inputs[0]->getGuid() - << ",B=" << inputs[1]->getGuid() << ",C=" << outputs[0]->getGuid() - << ",mnk=[" << m << "," << n << "," << k << "])"; - return os.str(); - } - - optional> MatmulObj::inferShape(const TensorVec &inputs) - { - // =================================== 作业 =================================== - // TODO：返回经过 matmul 操作后的 shape - // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm - // =================================== 作业 =================================== - return std::nullopt; - } - +#include "operators/matmul.h" + +namespace infini +{ + + MatmulObj::MatmulObj(GraphObj *graph, Tensor A, Tensor B, Tensor C, bool transA, + bool transB) + : OperatorObj(OpType::MatMul, TensorVec{A, B}, {C}), + transA(transA), transB(transB) + { + IT_ASSERT(checkValid(graph)); + } + + string MatmulObj::toString() const + { + std::ostringstream os; + os << "Matmul([" << (transA ? "A^T" : "A") << "," << (transB ? "B^T" : "B]") + << ",A=" << inputs[0]->getGuid() + << ",B=" << inputs[1]->getGuid() << ",C=" << outputs[0]->getGuid() + << ",mnk=[" << m << "," << n << "," << k << "])"; + return os.str(); + } + + optional> MatmulObj::inferShape(const TensorVec &inputs) + { + // =================================== 作业 =================================== + // TODO：返回经过 matmul 操作后的 shape + // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm + // =================================== 作业 =================================== + const auto A = inputs[0], B = inputs[1]; + auto shapeA = A->getDims(); + auto shapeB = B->getDims(); + int rankA = shapeA.size(); + int rankB = shapeB.size(); + + // Extract matrix dimensions + int aM = shapeA[rankA - 2], aK = shapeA[rankA - 1]; + int bK = shapeB[rankB - 2], bN = shapeB[rankB - 1]; + if (transA) + std::swap(aM, aK); + if (transB) + std::swap(bK, bN); + + m = aM; + n = bN; + k = aK; + + // Broadcast batch dimensions + int batchRank = std::max(rankA, rankB) - 2; + Shape outputShape(batchRank + 2); + for (int i = 0; i < batchRank; ++i) { + int dimA = (i < rankA - 2) ? shapeA[i] : 1; + int dimB = (i < rankB - 2) ? shapeB[i] : 1; + IT_ASSERT(dimA == dimB || dimA == 1 || dimB == 1, + "Incompatible batch dimensions for matmul"); + outputShape[i] = std::max(dimA, dimB); + } + outputShape[batchRank] = m; + outputShape[batchRank + 1] = n; + + return {{outputShape}}; + } + } // namespace infini \ No newline at end of file diff --git a/src/operators/transpose.cc b/src/operators/transpose.cc index faab2b69..d15d36af 100644 --- a/src/operators/transpose.cc +++ b/src/operators/transpose.cc @@ -1,50 +1,53 @@ -#include "operators/transpose.h" - -namespace infini -{ - TransposeObj::TransposeObj(GraphObj *graph, Tensor input, Tensor output, - vector permute) - : OperatorObj(OpType::Transpose, {input}, {output}) - { - auto rank = input->getRank(); - if (permute.empty()) - { - for (size_t i = 0; i < rank; ++i) - { - transposePermute[i] = i; - } - } - else - { - IT_ASSERT(rank == permute.size()); - transposePermute = std::move(permute); - } - IT_ASSERT(checkValid(graph)); - } - - optional> TransposeObj::inferShape(const TensorVec &inputs) - { - const auto A = inputs[0]; - auto input_dim = A->getDims(); - 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 - // =================================== 作业 =================================== - - return std::nullopt; - } - - std::string TransposeObj::toString() const - { - std::ostringstream os; - os << type.toString() << "[" << getGuid() << "]"; - os << "("; - os << vecToString(inputs[0]->getDims()) << ","; - os << "input=" << inputs[0]->getGuid() << ","; - os << "output=" << outputs[0]->getGuid() << ")"; - return os.str(); - } -}; // namespace infini +#include "operators/transpose.h" + +namespace infini +{ + TransposeObj::TransposeObj(GraphObj *graph, Tensor input, Tensor output, + vector permute) + : OperatorObj(OpType::Transpose, {input}, {output}) + { + auto rank = input->getRank(); + if (permute.empty()) + { + for (size_t i = 0; i < rank; ++i) + { + transposePermute[i] = i; + } + } + else + { + IT_ASSERT(rank == permute.size()); + transposePermute = std::move(permute); + } + IT_ASSERT(checkValid(graph)); + } + + optional> TransposeObj::inferShape(const TensorVec &inputs) + { + const auto A = inputs[0]; + auto input_dim = A->getDims(); + 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 + // =================================== 作业 =================================== + for (int i = 0; i < rank; ++i) { + output_dim[i] = input_dim[transposePermute[i]]; + } + + return {{output_dim}}; + } + + std::string TransposeObj::toString() const + { + std::ostringstream os; + os << type.toString() << "[" << getGuid() << "]"; + os << "("; + os << vecToString(inputs[0]->getDims()) << ","; + os << "input=" << inputs[0]->getGuid() << ","; + os << "output=" << outputs[0]->getGuid() << ")"; + return os.str(); + } +}; // namespace infini diff --git a/src/operators/unary.cc b/src/operators/unary.cc index 3daad361..b46db1a3 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -1,148 +1,148 @@ -#include "operators/unary.h" - -namespace infini -{ - UnaryObj::UnaryObj(OpType type, GraphObj *graph, Tensor input, Tensor output) - : OperatorObj(type, {input}, {output}) - { - IT_ASSERT(checkValid(graph)); - } - - optional> UnaryObj::inferShape(const TensorVec &inputs) - { - const auto A = inputs[0]; - return {{A->getDims()}}; - } - - std::string UnaryObj::toString() const - { - std::ostringstream os; - os << type.toString() << "[" << getGuid() << "]"; - os << "("; - os << vecToString(inputs[0]->getDims()) << ","; - os << "input=" << inputs[0]->getGuid() << ","; - os << "output=" << outputs[0]->getGuid() << ")"; - return os.str(); - } - - ClipObj::ClipObj(GraphObj *graph, Tensor input, Tensor output, - std::optional min, std::optional max) - : OperatorObj(OpType::Clip, {input}, {output}), minValue(min), - maxValue(max) - { - IT_ASSERT(checkValid(graph)); - } - - optional> ClipObj::inferShape(const TensorVec &inputs) - { - // =================================== 作业 =================================== - // TODO：返回经过 clip 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 - // =================================== 作业 =================================== - return std::nullopt; - } - - std::string ClipObj::toString() const - { - std::ostringstream os; - os << type.toString() << "[" << getGuid() << "]"; - os << "("; - os << vecToString(inputs[0]->getDims()) << ","; - os << "input=" << inputs[0]->getGuid() << ","; - os << "output=" << outputs[0]->getGuid() << ")"; - return os.str(); - } - - CastObj::CastObj(GraphObj *graph, Tensor input, Tensor output, CastType type) - : OperatorObj(OpType::Cast, {input}, {output}), castType(type) - { - IT_ASSERT(checkValid(graph)); - } - - 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 {}; - } - - optional> CastObj::inferShape(const TensorVec &inputs) - { - // =================================== 作业 =================================== - // TODO：返回经过 cast 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 - // =================================== 作业 =================================== - return std::nullopt; - } - - std::string CastObj::toString() const - { - std::ostringstream os; - os << type.toString() << "[" << getGuid() << "]"; - os << "("; - os << "output=" << outputs[0]->getGuid() << ")"; - return os.str(); - } - - DataType CastObj::getOutputDataType() const - { - switch (castType) - { - case CastType::Float2Float16: - return DataType::Float16; - case CastType::Float2Int64: - return DataType::Int64; - case CastType::Float2Int32: - return DataType::Int32; - case CastType::Float2Int16: - return DataType::Int16; - case CastType::Float2Int8: - return DataType::Int8; - case CastType::Int322Float: - return DataType::Float32; - case CastType::Int322Int8: - return DataType::Int8; - case CastType::Int322Int16: - return DataType::Int16; - case CastType::Int162Float: - return DataType::Float32; - case CastType::Int162Int32: - return DataType::Int32; - case CastType::Int82Float: - return DataType::Float32; - case CastType::Int82Int16: - return DataType::Int16; - case CastType::Int82Int32: - return DataType::Int32; - case CastType::Uint82Float: - return DataType::Float32; - case CastType::Uint82Int32: - return DataType::Int32; - case CastType::Uint82Int64: - return DataType::Int64; - case CastType::Int322Int64: - return DataType::Int64; - case CastType::Int642Int32: - return DataType::Int32; - case CastType::Int642Uint32: - return DataType::UInt32; - case CastType::Int642Float: - return DataType::Float32; - case CastType::Uint322Int64: - return DataType::Int64; - case CastType::Float162Float: - return DataType::Float32; - case CastType::BFloat162Float: - return DataType::Float32; - case CastType::Float2BFloat16: - return DataType::BFloat16; - case CastType::Float2Float: - return DataType::Float32; - default: - IT_TODO_HALT(); - } - } -}; // namespace infini +#include "operators/unary.h" + +namespace infini +{ + UnaryObj::UnaryObj(OpType type, GraphObj *graph, Tensor input, Tensor output) + : OperatorObj(type, {input}, {output}) + { + IT_ASSERT(checkValid(graph)); + } + + optional> UnaryObj::inferShape(const TensorVec &inputs) + { + const auto A = inputs[0]; + return {{A->getDims()}}; + } + + std::string UnaryObj::toString() const + { + std::ostringstream os; + os << type.toString() << "[" << getGuid() << "]"; + os << "("; + os << vecToString(inputs[0]->getDims()) << ","; + os << "input=" << inputs[0]->getGuid() << ","; + os << "output=" << outputs[0]->getGuid() << ")"; + return os.str(); + } + + ClipObj::ClipObj(GraphObj *graph, Tensor input, Tensor output, + std::optional min, std::optional max) + : OperatorObj(OpType::Clip, {input}, {output}), minValue(min), + maxValue(max) + { + IT_ASSERT(checkValid(graph)); + } + + optional> ClipObj::inferShape(const TensorVec &inputs) + { + // =================================== 作业 =================================== + // TODO：返回经过 clip 操作后的 shape + // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 + // =================================== 作业 =================================== + return {{inputs[0]->getDims()}}; + } + + std::string ClipObj::toString() const + { + std::ostringstream os; + os << type.toString() << "[" << getGuid() << "]"; + os << "("; + os << vecToString(inputs[0]->getDims()) << ","; + os << "input=" << inputs[0]->getGuid() << ","; + os << "output=" << outputs[0]->getGuid() << ")"; + return os.str(); + } + + CastObj::CastObj(GraphObj *graph, Tensor input, Tensor output, CastType type) + : OperatorObj(OpType::Cast, {input}, {output}), castType(type) + { + IT_ASSERT(checkValid(graph)); + } + + 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 {getOutputDataType()}; + } + + optional> CastObj::inferShape(const TensorVec &inputs) + { + // =================================== 作业 =================================== + // TODO：返回经过 cast 操作后的 shape + // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 + // =================================== 作业 =================================== + return {{inputs[0]->getDims()}}; + } + + std::string CastObj::toString() const + { + std::ostringstream os; + os << type.toString() << "[" << getGuid() << "]"; + os << "("; + os << "output=" << outputs[0]->getGuid() << ")"; + return os.str(); + } + + DataType CastObj::getOutputDataType() const + { + switch (castType) + { + case CastType::Float2Float16: + return DataType::Float16; + case CastType::Float2Int64: + return DataType::Int64; + case CastType::Float2Int32: + return DataType::Int32; + case CastType::Float2Int16: + return DataType::Int16; + case CastType::Float2Int8: + return DataType::Int8; + case CastType::Int322Float: + return DataType::Float32; + case CastType::Int322Int8: + return DataType::Int8; + case CastType::Int322Int16: + return DataType::Int16; + case CastType::Int162Float: + return DataType::Float32; + case CastType::Int162Int32: + return DataType::Int32; + case CastType::Int82Float: + return DataType::Float32; + case CastType::Int82Int16: + return DataType::Int16; + case CastType::Int82Int32: + return DataType::Int32; + case CastType::Uint82Float: + return DataType::Float32; + case CastType::Uint82Int32: + return DataType::Int32; + case CastType::Uint82Int64: + return DataType::Int64; + case CastType::Int322Int64: + return DataType::Int64; + case CastType::Int642Int32: + return DataType::Int32; + case CastType::Int642Uint32: + return DataType::UInt32; + case CastType::Int642Float: + return DataType::Float32; + case CastType::Uint322Int64: + return DataType::Int64; + case CastType::Float162Float: + return DataType::Float32; + case CastType::BFloat162Float: + return DataType::Float32; + case CastType::Float2BFloat16: + return DataType::BFloat16; + case CastType::Float2Float: + return DataType::Float32; + default: + IT_TODO_HALT(); + } + } +}; // namespace infini diff --git a/src/utils/exception.cc b/src/utils/exception.cc index 228a39a8..c1d3814d 100644 --- a/src/utils/exception.cc +++ b/src/utils/exception.cc @@ -1,5 +1,5 @@ -#include "utils/exception.h" - -namespace infini { -Exception::Exception(const std::string &msg) : std::runtime_error(msg) {} -} // namespace infini +#include "utils/exception.h" + +namespace infini { +Exception::Exception(const std::string &msg) : std::runtime_error(msg) {} +} // namespace infini diff --git a/src/utils/operator_utils.cc b/src/utils/operator_utils.cc index edbd2c82..e68ee4ea 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -1,69 +1,80 @@ -#include "utils/operator_utils.h" -#include "core/runtime.h" - -namespace infini { - -Shape infer_broadcast(const Shape &A, const Shape &B) { - - // =================================== 作业 =================================== - // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 - // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md - // =================================== 作业 =================================== - - return {}; -} - -int get_real_axis(const int &axis, const int &rank) { - IT_ASSERT(rank >= 1); - IT_ASSERT(axis >= -rank && axis <= (rank - 1)); - int newAxis; - if (axis < 0) { - newAxis = rank + axis; - } else { - newAxis = axis; - } - return newAxis; -} - -Shape locate_index(size_t inputN, const Shape &shape) { - Shape ans(shape.size()); - auto i = ans.rbegin(); - auto j = shape.rbegin(), ej = shape.rend(); - while (j != ej) { - auto div = std::div(inputN, *j++); - *i++ = div.rem; - inputN = div.quot; - } - return ans; -} - -size_t delocate_index(const Shape &shapeIndex, const Shape &shape, - const Shape &stride) { - size_t ans = 0; - Shape index(shapeIndex.size()); - IT_ASSERT(shapeIndex.size() == shape.size()); - IT_ASSERT(shape.size() == stride.size()); - for (size_t i = 0; i < shape.size(); ++i) { - index[i] = shapeIndex[i] % shape[i]; - ans += index[i] * stride[i]; - } - return ans; -} - -std::string device_to_str(Device device) { - std::string deviceStr; - switch (device) { - case Device::CPU: - return "CPU"; - default: - IT_TODO_HALT(); - } -} - -std::string get_kernel_attrs_str(const KernelAttrs &kernelAttrs) { - std::string deviceStr = device_to_str(std::get<0>(kernelAttrs)); - std::string opStr = OpType(std::get<1>(kernelAttrs)).toString(); - return deviceStr + ", " + opStr; -} - -} // namespace infini +#include "utils/operator_utils.h" +#include "core/runtime.h" + +namespace infini { + +Shape infer_broadcast(const Shape &A, const Shape &B) { + + // =================================== 作业 =================================== + // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 + // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md + // =================================== 作业 =================================== + + auto rankA = A.size(); + auto rankB = B.size(); + auto rank = std::max(rankA, rankB); + Shape result(rank); + for (size_t i = 0; i < rank; ++i) { + auto dimA = (i < rank - rankA) ? 1 : A[i - (rank - rankA)]; + auto dimB = (i < rank - rankB) ? 1 : B[i - (rank - rankB)]; + IT_ASSERT(dimA == dimB || dimA == 1 || dimB == 1, + "Incompatible broadcast shapes"); + result[i] = std::max(dimA, dimB); + } + return result; +} + +int get_real_axis(const int &axis, const int &rank) { + IT_ASSERT(rank >= 1); + IT_ASSERT(axis >= -rank && axis <= (rank - 1)); + int newAxis; + if (axis < 0) { + newAxis = rank + axis; + } else { + newAxis = axis; + } + return newAxis; +} + +Shape locate_index(size_t inputN, const Shape &shape) { + Shape ans(shape.size()); + auto i = ans.rbegin(); + auto j = shape.rbegin(), ej = shape.rend(); + while (j != ej) { + auto div = std::div(inputN, *j++); + *i++ = div.rem; + inputN = div.quot; + } + return ans; +} + +size_t delocate_index(const Shape &shapeIndex, const Shape &shape, + const Shape &stride) { + size_t ans = 0; + Shape index(shapeIndex.size()); + IT_ASSERT(shapeIndex.size() == shape.size()); + IT_ASSERT(shape.size() == stride.size()); + for (size_t i = 0; i < shape.size(); ++i) { + index[i] = shapeIndex[i] % shape[i]; + ans += index[i] * stride[i]; + } + return ans; +} + +std::string device_to_str(Device device) { + std::string deviceStr; + switch (device) { + case Device::CPU: + return "CPU"; + default: + IT_TODO_HALT(); + } +} + +std::string get_kernel_attrs_str(const KernelAttrs &kernelAttrs) { + std::string deviceStr = device_to_str(std::get<0>(kernelAttrs)); + std::string opStr = OpType(std::get<1>(kernelAttrs)).toString(); + return deviceStr + ", " + opStr; +} + +} // namespace infini diff --git "a/\346\210\252\345\233\276.png" "b/\346\210\252\345\233\276.png" new file mode 100644 index 00000000..41e815be Binary files /dev/null and "b/\346\210\252\345\233\276.png" differ