/*! * Copyright (c) 2016 by Contributors * Implementation of API functions related to IR build * \file c_api_ir.cc */ #include <tvm/expr.h> #include <tvm/ir.h> #include <ir/IROperator.h> #include "./c_api_registry.h" namespace tvm { namespace ir { using ArgStack = const std::vector<APIVariantValue>; using RetValue = APIVariantValue; TVM_REGISTER_API(_Var) .set_body([](const ArgStack& args, RetValue *ret) { *ret = Variable::make(args.at(1), args.at(0)); }); TVM_REGISTER_API(_make_For) .set_body([](const ArgStack& args, RetValue *ret) { *ret = For::make(args.at(0), args.at(1), args.at(2), static_cast<ForType>(args.at(3).operator int()), static_cast<Halide::DeviceAPI>(args.at(4).operator int()), args.at(5)); }); TVM_REGISTER_API(_make_Call) .set_body([](const ArgStack& args, RetValue *ret) { *ret = Call::make(args.at(0), args.at(1), args.at(2), static_cast<Call::CallType>(args.at(3).operator int()), args.at(4)); }); TVM_REGISTER_API(_make_Allocate) .set_body([](const ArgStack& args, RetValue *ret) { *ret = Allocate::make(args.at(0), args.at(1), args.at(2), args.at(3), args.at(4)); }); // make from two arguments #define REGISTER_MAKE1(Node) \ TVM_REGISTER_API(_make_## Node) \ .set_body([](const ArgStack& args, RetValue *ret) { \ *ret = Node::make(args.at(0)); \ }) \ #define REGISTER_MAKE2(Node) \ TVM_REGISTER_API(_make_## Node) \ .set_body([](const ArgStack& args, RetValue *ret) { \ *ret = Node::make(args.at(0), args.at(1)); \ }) \ #define REGISTER_MAKE3(Node) \ TVM_REGISTER_API(_make_## Node) \ .set_body([](const ArgStack& args, RetValue *ret) { \ *ret = Node::make(args.at(0), args.at(1), args.at(2)); \ }) \ #define REGISTER_MAKE4(Node) \ TVM_REGISTER_API(_make_## Node) \ .set_body([](const ArgStack& args, RetValue *ret) { \ *ret = Node::make(args.at(0), args.at(1), args.at(2), args.at(3)); \ }) \ #define REGISTER_MAKE_BINARY_OP(Node) \ TVM_REGISTER_API(_make_## Node) \ .set_body([](const ArgStack& args, RetValue *ret) { \ Expr a = args.at(0), b = args.at(1); \ match_types(a, b); \ *ret = Node::make(a, b); \ }) \ .add_argument("lhs", "Expr", "left operand") \ .add_argument("rhs", "Expr", "right operand") REGISTER_MAKE3(Reduce); REGISTER_MAKE4(AttrStmt); REGISTER_MAKE2(IntImm); REGISTER_MAKE2(UIntImm); REGISTER_MAKE2(FloatImm); REGISTER_MAKE1(StringImm); REGISTER_MAKE_BINARY_OP(Add); REGISTER_MAKE_BINARY_OP(Sub); REGISTER_MAKE_BINARY_OP(Mul); REGISTER_MAKE_BINARY_OP(Div); REGISTER_MAKE_BINARY_OP(Mod); REGISTER_MAKE_BINARY_OP(Min); REGISTER_MAKE_BINARY_OP(Max); REGISTER_MAKE_BINARY_OP(EQ); REGISTER_MAKE_BINARY_OP(NE); REGISTER_MAKE_BINARY_OP(LT); REGISTER_MAKE_BINARY_OP(LE); REGISTER_MAKE_BINARY_OP(GT); REGISTER_MAKE_BINARY_OP(GE); REGISTER_MAKE_BINARY_OP(And); REGISTER_MAKE_BINARY_OP(Or); REGISTER_MAKE1(Not); REGISTER_MAKE3(Select); REGISTER_MAKE3(Ramp); REGISTER_MAKE2(Broadcast); REGISTER_MAKE3(Let); REGISTER_MAKE3(LetStmt); REGISTER_MAKE2(AssertStmt); REGISTER_MAKE3(ProducerConsumer); REGISTER_MAKE3(Store); REGISTER_MAKE3(Provide); REGISTER_MAKE1(Free); // TODO(tqchen) Realize; REGISTER_MAKE2(Block); REGISTER_MAKE3(IfThenElse); REGISTER_MAKE1(Evaluate); } // namespace ir } // namespace tvm