mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-29 23:06:41 +00:00
0d204f3f06
Description: This change adds the common part of TVM based codegen library. It includes following parts: * Microsoft TVM Inventory (MTI): a set of TVM ops for neural networks, similar to TOPI * Compiler pass for traversing ONNX graph and generate TVM ops * Compiler pass for traversing generated graph and specify TVM schedule * Compiler pass for handling weight layout * Utils for debugging Motivation and Context: TVM is an open deep learning compiler stack for cpu, gpu and specialized accelerators. To leverage it in ONNX, we built an execution provider named Nuphar. Currently, Nuphar gets good performance on CPUs with AVX2 on quantized LSTM models. This codegen library was part of Nuphar execution provider. It is split out for sharing with other execution providers, as we'd like to reuse TVM in more devices.
18 lines
585 B
C++
18 lines
585 B
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/codegen/mti/math/logsoftmax.h"
|
|
|
|
#include "core/codegen/mti/tensor/reshape_ops.h"
|
|
#include <topi/nn/softmax.h>
|
|
|
|
namespace onnxruntime {
|
|
namespace tvm_codegen {
|
|
|
|
tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, const std::string& name) {
|
|
tvm::Tensor flatten_t = Flatten(input, axis, "logsoftmax_flatten");
|
|
return Reshape(topi::nn::log_softmax(flatten_t, name), input->shape, "logsoftmax_reshape");
|
|
}
|
|
|
|
} // namespace tvm_codegen
|
|
} // namespace onnxruntime
|