Skip to content

Instantly share code, notes, and snippets.

@sergei-mironov
Created December 4, 2018 20:05
Show Gist options
  • Save sergei-mironov/fafa70cde68cfc33b11c36bd20b3e6f0 to your computer and use it in GitHub Desktop.
Save sergei-mironov/fafa70cde68cfc33b11c36bd20b3e6f0 to your computer and use it in GitHub Desktop.
Sigmoid problem illustration
/*
g++ -std=c++14 sigmoid.cpp -ltvm -o sigmoid.gen
sigmoid.gen > sigmoid.s
*/
#include <random>
#include <iomanip>
#include <array>
#include <exception>
#include <tvm/tvm.h>
#include <tvm/operation.h>
#include <tvm/schedule.h>
#include <tvm/tensor.h>
#include <tvm/build_module.h>
#include <tvm/ir.h>
#include <tvm/ir_operator.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_pass.h>
#include <topi/broadcast.h>
#include <topi/reduction.h>
#include <ir/IR.h>
#include <ir/IRPrinter.h>
using namespace std;
using namespace tvm;
int main()
{
BuildConfig config = build_config();
auto n = var("n");
Array<Expr> shape = {n,n};
Tensor A = placeholder(shape, Float(32), "A");
Tensor X = compute(shape, FCompute([=](auto i){ return tvm::sigmoid(A(i)); } )) ;
auto vecadd_lowered = ({
Schedule s = create_schedule({X->op});
std::unordered_map<Tensor, Buffer> binds;
auto args = Array<Tensor>({A, X});
auto lowered = lower(s, args, "vecadd", binds, config);
lowered;
});
auto target = Target::create("llvm");
auto target_host = Target::create("llvm");
runtime::Module mod = build(vecadd_lowered, target, target_host, config);
/* Output LLVM assembly to stdout */
cout << mod->GetSource("asm") << endl;
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment