| #include "caffe2/core/net.h" |
| #include "caffe2/utils/proto_utils.h" |
| |
| #include "compiler.h" |
| #include "parser.h" |
| |
| namespace caffe2 { |
| namespace script { |
| |
| namespace { |
| |
| static std::unordered_set<std::string> ops_containing_nets = { |
| "If", |
| "While", |
| "RecurrentNetwork", |
| }; |
| // record of defined function |
| // NetDef + metadata |
| struct FunctionDefinition { |
| explicit FunctionDefinition(Def tree) |
| : tree(new Def(tree)), net_def(new NetDef()) {} |
| |
| explicit FunctionDefinition(std::unique_ptr<NetDef> def) |
| : tree(nullptr), net_def(std::move(def)) { |
| // we coop extern_inputs/extern_outputs to be the inputs/outputs to |
| // this net as a function |
| // but we _dont_ set these when creating the net in the workspace |
| // because they require the net to have valid inputs/outputs |
| inputs.insert( |
| inputs.begin(), |
| net_def->external_input().begin(), |
| net_def->external_input().end()); |
| outputs.insert( |
| outputs.begin(), |
| net_def->external_output().begin(), |
| net_def->external_output().end()); |
| net_def->clear_external_output(); |
| net_def->clear_external_input(); |
| } |
| |
| bool isExtern() const { |
| return tree == nullptr; |
| } |
| std::unique_ptr<Def> tree; |
| std::unique_ptr<NetDef> net_def; |
| std::vector<std::string> inputs; |
| std::vector<std::string> outputs; |
| }; |
| |
| } // namespace |
| |
| using SymbolTable = std::unordered_map<std::string, FunctionDefinition>; |
| |
| struct DefCompiler { |
| DefCompiler(FunctionDefinition& def, SymbolTable& symbol_table) |
| : def(def), |
| net_def_stack({def.net_def.get()}), |
| symbol_table(symbol_table) {} |
| void run() { |
| auto& tree = *def.tree; |
| cur().set_name(tree.name().name()); |
| for (auto input : tree.params()) { |
| auto& name = input.ident().name(); |
| map(name, name); |
| def.inputs.push_back(name); |
| } |
| for (auto output : tree.returns()) { |
| auto& name = output.ident().name(); |
| map(name, name); |
| def.outputs.push_back(name); |
| } |
| emitStatements(tree.statements()); |
| } |
| void emitExpressionStatement(TreeRef stmt) { |
| // expression with no used outputs |
| emit(stmt, {}); |
| } |
| void emitStatements(const ListView<TreeRef>& statements) { |
| for (auto stmt : statements) { |
| switch (stmt->kind()) { |
| case TK_IF: |
| emitIf(If(stmt)); |
| break; |
| case TK_WHILE: |
| emitWhile(While(stmt)); |
| break; |
| case TK_ASSIGN: |
| emitAssignment(Assign(stmt)); |
| break; |
| case TK_GLOBAL: |
| for (auto ident : stmt->trees()) { |
| auto name = Ident(ident).name(); |
| map(name, name); |
| } |
| break; |
| default: |
| emitExpressionStatement(stmt); |
| break; |
| } |
| } |
| } |
| void map(const std::string& name, const std::string& value) { |
| env[name] = value; |
| } |
| const std::string& lookup(const Ident& ident) { |
| if (env.count(ident.name()) == 0) |
| throw ErrorReport(ident) << "undefined value " << ident.name(); |
| return env[ident.name()]; |
| } |
| void emitAssignment(const Assign& stmt) { |
| std::vector<std::string> outputs; |
| for (auto lhs : stmt.lhs()) { |
| std::string name = getLHS(lhs); |
| // use of "_" gets renamed in Caffe2 graphs so that two uses |
| // don't unintentionally interfere with each other |
| if (name == "_") { |
| name = fresh(); |
| } |
| outputs.push_back(name); |
| } |
| if (stmt.reduction() != '=') { |
| if (stmt.lhs().size() != 1) { |
| throw ErrorReport(stmt) |
| << "reductions are only allow when there is a single variable " |
| << "on the left-hand side."; |
| } |
| auto lhs = stmt.lhs()[0]; |
| auto expr = |
| Compound::create(stmt.reduction(), stmt.range(), {lhs, stmt.rhs()}); |
| emit(expr, outputs); |
| } else { |
| emit(stmt.rhs(), outputs); |
| } |
| int i = 0; |
| for (auto ident : stmt.lhs()) { |
| if (ident->kind() == TK_IDENT) |
| map(Ident(ident).name(), outputs.at(i)); |
| i++; |
| } |
| } |
| void emitIf(const If& stmt) { |
| auto cond = getValue(stmt.cond()); |
| auto op = cur().add_op(); |
| op->set_type("If"); |
| op->add_input(cond); |
| auto true_branch = op->add_arg(); |
| true_branch->set_name("then_net"); |
| auto nd = true_branch->mutable_n(); |
| net_def_stack.push_back(nd); |
| emitStatements(stmt.trueBranch()); |
| net_def_stack.pop_back(); |
| if (stmt.falseBranch().size() > 0) { |
| auto false_branch = op->add_arg(); |
| false_branch->set_name("else_net"); |
| auto nd = false_branch->mutable_n(); |
| net_def_stack.push_back(nd); |
| emitStatements(stmt.falseBranch()); |
| net_def_stack.pop_back(); |
| } |
| } |
| void emitWhile(const While& stmt) { |
| std::string loop_var = fresh(); |
| emitConst(0, loop_var, "i"); // it needs a definition before loop |
| auto op = cur().add_op(); |
| op->set_type("While"); |
| auto cond = op->add_arg(); |
| cond->set_name("cond_net"); |
| auto cond_net = cond->mutable_n(); |
| |
| net_def_stack.push_back(cond_net); |
| emit(stmt.cond(), {loop_var}); |
| net_def_stack.pop_back(); |
| |
| op->add_input(loop_var); |
| auto body = op->add_arg(); |
| body->set_name("loop_net"); |
| auto body_net = body->mutable_n(); |
| |
| net_def_stack.push_back(body_net); |
| emitStatements(stmt.body()); |
| net_def_stack.pop_back(); |
| } |
| std::string getLHS(const TreeRef& tree) { |
| switch (tree->kind()) { |
| case TK_IDENT: { |
| return Ident(tree).name(); |
| } break; |
| case '.': { |
| auto sel = Select(tree); |
| std::string lhs = getValue(sel.value()); |
| // TODO: check whether this subname exists in object lhs |
| return lhs + "/" + sel.selector().name(); |
| } break; |
| default: { |
| throw ErrorReport(tree) |
| << "This expression cannot appear on the left-hand size of an assignment"; |
| } break; |
| } |
| } |
| std::string getValue(const TreeRef& tree) { |
| switch (tree->kind()) { |
| case TK_IDENT: { |
| return lookup(Ident(tree)); |
| } break; |
| case '.': { |
| auto sel = Select(tree); |
| std::string lhs = getValue(sel.value()); |
| // TODO: check whether this subname exists in object lhs |
| return lhs + "/" + sel.selector().name(); |
| } break; |
| default: { |
| std::string name = fresh(); |
| emit(tree, {name}); |
| return name; |
| } break; |
| } |
| } |
| std::string fresh(std::string prefix = "$t") { |
| return std::string(prefix) + c10::to_string(next_fresh++); |
| } |
| const char* operatorName(int kind, int ninputs) { |
| switch (kind) { |
| case '+': |
| return "Add"; |
| case '-': |
| if (ninputs == 1) |
| return "Negative"; |
| else |
| return "Sub"; |
| case '*': |
| return "Mul"; |
| case '/': |
| return "Div"; |
| case TK_NE: |
| return "NE"; |
| case TK_EQ: |
| return "EQ"; |
| case '<': |
| return "LT"; |
| case '>': |
| return "GT"; |
| case TK_LE: |
| return "LE"; |
| case TK_GE: |
| return "GE"; |
| case TK_IF_EXPR: |
| return "Conditional"; |
| case TK_AND: |
| return "And"; |
| case TK_OR: |
| return "Or"; |
| case TK_NOT: |
| return "Not"; |
| default: |
| throw std::runtime_error("unknown kind " + c10::to_string(kind)); |
| } |
| } |
| void fillArg(Argument* arg, const Attribute& attr) { |
| std::string name = attr.name().name(); |
| arg->set_name(name); |
| auto value = attr.value(); |
| // TODO: handle non-float attributes |
| switch (value->kind()) { |
| case TK_CONST: { |
| auto v = value->tree(0)->doubleValue(); |
| auto f = value->tree(1)->stringValue(); |
| if (f == "f") |
| arg->set_f(v); |
| else |
| arg->set_i(v); |
| } break; |
| case TK_LIST: |
| for (auto t : value->trees()) { |
| auto v = t->tree(0)->doubleValue(); |
| auto f = t->tree(1)->stringValue(); |
| if (f == "f") |
| arg->add_floats(v); |
| else |
| arg->add_ints(v); |
| } |
| break; |
| } |
| } |
| template <typename Trees> |
| std::vector<std::string> getValues(const Trees& trees) { |
| std::vector<std::string> result; |
| for (const auto& tree : trees) { |
| result.push_back(getValue(tree)); |
| } |
| return result; |
| } |
| |
| bool renameLookup( |
| std::unordered_map<std::string, std::string>& rename_map, |
| const std::string& name, |
| std::string& rename) { |
| // first look for name in the map directly |
| auto it = rename_map.find(name); |
| if (it != rename_map.end()) { |
| rename = it->second; |
| return true; |
| } |
| // otherwise if we have a rename entry like a => b and a name "a/foo/bar" |
| // then replace it with "b/foo/bar" |
| auto p = name.find("/"); |
| if (p == std::string::npos) |
| return false; |
| it = rename_map.find(name.substr(0, p)); |
| if (it != rename_map.end()) { |
| rename = it->second + name.substr(p); |
| return true; |
| } |
| return false; |
| } |
| void renameOp( |
| std::unordered_map<std::string, std::string>& rename_map, |
| const Apply& apply, |
| const std::string& prefix, |
| bool isExtern, |
| OperatorDef* new_op) { |
| for (size_t i = 0; i < new_op->input().size(); i++) { |
| auto& name = new_op->input(i); |
| std::string renamed; |
| bool defined = renameLookup(rename_map, name, renamed); |
| if (!isExtern && !defined) { |
| throw ErrorReport(apply) |
| << " unexpected undefined name '" << name |
| << "' while attempting to inline '" << apply.name().name() << "'"; |
| } else if (!defined) { |
| // extern function using a global name, assign it an identity mapping |
| rename_map[name] = name; |
| } |
| new_op->set_input(i, renamed); |
| } |
| for (size_t i = 0; i < new_op->output().size(); i++) { |
| auto& name = new_op->output(i); |
| std::string renamed; |
| if (!renameLookup(rename_map, name, renamed)) { |
| renamed = prefix + name; |
| rename_map[name] = renamed; |
| } |
| new_op->set_output(i, renamed); |
| } |
| // handle control flow inside the op as well |
| if (ops_containing_nets.count(new_op->type()) > 0) { |
| for (size_t i = 0; i < new_op->arg_size(); i++) { |
| auto* arg = new_op->mutable_arg(i); |
| if (arg->has_n()) { |
| auto* n = arg->mutable_n(); |
| for (size_t j = 0; j < n->op_size(); j++) { |
| renameOp(rename_map, apply, prefix, isExtern, n->mutable_op(j)); |
| } |
| } |
| } |
| } |
| } |
| |
| bool hasBypassRename(const Apply& apply) { |
| for (auto attr : apply.attributes()) { |
| if (attr.name().name() == "rename") { |
| if (attr.value()->kind() != TK_CONST) { |
| throw ErrorReport(attr.value()) << "expected a single constant"; |
| } |
| return attr.value()->tree(0)->doubleValue() == 0; |
| } |
| } |
| return false; |
| } |
| |
| // emit a function call by inlining the function's NetDef into our |
| // net def, renaming temporaries func_name<unique_id>/orig_name |
| // renaming only happens for values defined by the function |
| // that are not marked outputs |
| |
| // inputs/outputs are passed by reference |
| void emitFunctionCall(Apply& apply, const std::vector<std::string>& outputs) { |
| std::string fname = apply.name().name(); |
| std::string prefix = fresh(fname) + "/"; |
| auto& fn = symbol_table.at(apply.name().name()); |
| bool isExtern = fn.isExtern(); |
| auto inputs = getValues(apply.inputs()); |
| std::unordered_map<std::string, std::string> rename_map; |
| if (inputs.size() != fn.inputs.size()) { |
| throw ErrorReport(apply) << fname << " expected " << fn.inputs.size() |
| << " values but received " << inputs.size(); |
| } |
| for (size_t i = 0; i < inputs.size(); i++) { |
| rename_map[fn.inputs[i]] = inputs[i]; |
| } |
| if (outputs.size() != fn.outputs.size()) { |
| throw ErrorReport(apply) << fname << " expected " << fn.outputs.size() |
| << " values but received " << outputs.size(); |
| } |
| for (size_t i = 0; i < outputs.size(); i++) { |
| rename_map[fn.outputs[i]] = outputs[i]; |
| } |
| for (auto& op : fn.net_def->op()) { |
| auto new_op = cur().add_op(); |
| new_op->CopyFrom(op); |
| if (hasBypassRename(apply)) { |
| prefix = ""; |
| } |
| renameOp(rename_map, apply, prefix, isExtern, new_op); |
| } |
| } |
| void expectOutputs( |
| const TreeRef& tree, |
| const std::vector<std::string>& outputs, |
| size_t size) { |
| if (outputs.size() != size) { |
| throw ErrorReport(tree) |
| << "expected operator to produce " << outputs.size() |
| << " outputs but it produced " << size; |
| } |
| } |
| void appendOutputs( |
| const TreeRef& tree, |
| OperatorDef* op, |
| const std::vector<std::string>& outputs, |
| size_t size) { |
| expectOutputs(tree, outputs, size); |
| for (size_t i = 0; i < size; i++) { |
| op->add_output(outputs[i]); |
| } |
| } |
| void emitOperator( |
| const Apply& apply, |
| const OpSchema* schema, |
| const std::vector<std::string>& outputs) { |
| // must be before add_op |
| auto values = getValues(apply.inputs()); |
| if (values.size() < schema->min_input() || |
| values.size() > schema->max_input()) { |
| if (schema->min_input() == schema->max_input()) { |
| throw ErrorReport(apply) << "operator expects " << schema->min_input() |
| << " inputs but found " << values.size(); |
| } else { |
| throw ErrorReport(apply) |
| << "operator takes between " << schema->min_input() << " and " |
| << schema->max_input() << " inputs but found " << values.size() |
| << "."; |
| } |
| } |
| auto numActualOutputs = schema->CalculateOutput(values.size()); |
| if (numActualOutputs != kCannotComputeNumOutputs && |
| outputs.size() != numActualOutputs) { |
| throw ErrorReport(apply) |
| << "operator produces " << numActualOutputs |
| << " outputs but matched to " << outputs.size() << " outputs"; |
| } |
| auto op = cur().add_op(); |
| op->set_type(apply.name().name()); |
| for (auto& v : values) { |
| op->add_input(v); |
| } |
| // assume 1 output unless matched to more |
| appendOutputs(apply, op, outputs, outputs.size()); |
| for (auto attribute : apply.attributes()) { |
| fillArg(op->add_arg(), attribute); |
| } |
| // Ok, we checked the stuff where we can easily give a friendly error |
| // message, now verify against the schema and report the error at the line |
| if (!schema->Verify(*op)) { |
| throw ErrorReport(apply) << "failed schema checking"; |
| } |
| } |
| |
| // Emit an operation, writing results into 'outputs'. |
| // This will _always_ compute something, unlike 'getValue' which simply |
| // returns an already computed reference if possible. |
| // So if 'tree' is an identifier or nested identifier (foo.bar) |
| // this will cause it to be _copied_ into outputs. |
| void emit(const TreeRef& tree, const std::vector<std::string>& outputs) { |
| switch (tree->kind()) { |
| case TK_IDENT: |
| case '.': { |
| auto op = cur().add_op(); |
| op->set_type("Copy"); |
| op->add_input(getValue(tree)); |
| appendOutputs(tree, op, outputs, 1); |
| } break; |
| case TK_NE: |
| case TK_EQ: |
| case '<': |
| case '>': |
| case TK_LE: |
| case TK_GE: |
| case '-': |
| case '*': |
| case '/': |
| case '+': |
| case TK_AND: |
| case TK_OR: |
| case TK_NOT: |
| case TK_IF_EXPR: { |
| // must be before add_op |
| auto values = getValues(tree->trees()); |
| auto op = cur().add_op(); |
| op->set_type(operatorName(tree->kind(), tree->trees().size())); |
| for (auto& v : values) { |
| op->add_input(v); |
| } |
| appendOutputs(tree, op, outputs, 1); |
| auto broadcast = op->add_arg(); |
| broadcast->set_name("broadcast"); |
| broadcast->set_i(1); |
| } break; |
| case TK_APPLY: { |
| auto apply = Apply(tree); |
| // Handle built-ins like zeros, ones, etc |
| if (builtins.count(apply.name().name()) > 0) { |
| builtins[apply.name().name()](this, apply, outputs); |
| break; |
| } |
| if (symbol_table.count(apply.name().name()) > 0) { |
| emitFunctionCall(apply, outputs); |
| break; |
| } |
| auto schema = OpSchemaRegistry::Schema(apply.name().name()); |
| if (schema) { |
| emitOperator(apply, schema, outputs); |
| break; |
| } |
| throw ErrorReport(apply) |
| << "attempting to call unknown operation or function '" |
| << apply.name().name() << "'"; |
| } break; |
| case TK_CAST: { |
| auto cast = Cast(tree); |
| auto c2type = getType(cast.type()); |
| auto input = getValue(cast.input()); |
| auto op = cur().add_op(); |
| op->set_type("Cast"); |
| op->add_input(input); |
| appendOutputs(tree, op, outputs, 1); |
| auto arg = op->add_arg(); |
| arg->set_name("to"); |
| arg->set_i(c2type); |
| } break; |
| case TK_CONST: { |
| expectOutputs(tree, outputs, 1); |
| emitConst( |
| tree->tree(0)->doubleValue(), |
| outputs[0], |
| tree->tree(1)->stringValue()); |
| } break; |
| case TK_GATHER: { |
| const auto gather = Gather(tree); |
| desugarAndEmitOperator( |
| "Gather", |
| gather.range(), |
| {gather.value(), gather.indices()}, |
| outputs); |
| break; |
| } |
| case TK_SLICE: { |
| const auto slice = Slice(tree); |
| desugarAndEmitOperator( |
| "Slice", |
| slice.range(), |
| {slice.value(), slice.startOr(0), slice.endOr(-1)}, |
| outputs); |
| break; |
| } |
| default: |
| throw ErrorReport(tree) << "NYI: " << tree; |
| break; |
| } |
| } |
| |
| // Desugars constructs that are syntactic sugar and emits the corresponding |
| // operator invocation, e.g. tensor[indices] -> tensor.Gather(indices). |
| void desugarAndEmitOperator( |
| const std::string& operatorName, |
| const SourceRange& range, |
| TreeList&& inputs, |
| const std::vector<std::string>& outputs) { |
| const auto applyName = Ident::create(range, operatorName); |
| const auto applyInputs = |
| Compound::create(TK_LIST, range, std::move(inputs)); |
| const auto applyAttributes = Compound::create(TK_LIST, range, {}); |
| const auto apply = |
| Apply::create(range, applyName, applyInputs, applyAttributes); |
| const auto schema = OpSchemaRegistry::Schema(operatorName); |
| assert(schema != nullptr); |
| emitOperator(Apply(apply), schema, outputs); |
| } |
| |
| TensorProto_DataType getType(int type) { |
| switch (type) { |
| case TK_INT: |
| return TensorProto_DataType_INT32; |
| case TK_FLOAT: |
| return TensorProto_DataType_FLOAT; |
| case TK_LONG: |
| return TensorProto_DataType_INT64; |
| case TK_BOOL: |
| return TensorProto_DataType_BOOL; |
| default: |
| throw std::runtime_error( |
| "expected type token: " + c10::to_string(type)); |
| } |
| } |
| |
| OperatorDef* emitConst( |
| double v, |
| const std::string& output, |
| const std::string& type_ident) { |
| auto op = cur().add_op(); |
| op->set_type("ConstantFill"); |
| auto dtype = op->add_arg(); |
| dtype->set_name("dtype"); |
| auto value = op->add_arg(); |
| value->set_name("value"); |
| if (type_ident == "f") { |
| dtype->set_i(TensorProto_DataType_FLOAT); |
| value->set_f(v); |
| } else if (type_ident == "LL") { |
| dtype->set_i(TensorProto_DataType_INT64); |
| value->set_i(v); |
| } else if (type_ident == "b") { |
| dtype->set_i(TensorProto_DataType_BOOL); |
| value->set_i(v != 0); |
| } else if (type_ident == "i") { |
| dtype->set_i(TensorProto_DataType_INT32); |
| value->set_i(v); |
| } else { |
| throw std::runtime_error("unknown type_ident " + type_ident); |
| } |
| auto shape = op->add_arg(); |
| shape->set_name("shape"); |
| shape->add_ints(1); |
| op->add_output(output); |
| return op; |
| } |
| NetDef& cur() { |
| return *net_def_stack.back(); |
| } |
| FunctionDefinition& def; // the def being constructed |
| std::unordered_map<std::string, std::string> |
| env; // map from name in Def to name in NetDef |
| std::vector<NetDef*> net_def_stack; |
| SymbolTable& symbol_table; |
| int next_fresh = 0; |
| |
| private: |
| void emitFillOp(const Apply& apply, const std::vector<std::string>& outputs) { |
| auto builtin_type = apply.name().name(); |
| auto values = getValues(apply.inputs()); |
| if (values.size() > 1) { |
| throw ErrorReport(apply) |
| << "Built-in " << builtin_type << " accepts 0 or 1 inputs."; |
| } |
| bool has_shape = false; |
| for (const auto& attribute : apply.attributes()) { |
| if (attribute.name().name() == "shape") { |
| has_shape = true; |
| } else { |
| throw ErrorReport(apply) |
| << "Unrecognized attribute " << attribute.name().name() |
| << " for built-in " << builtin_type; |
| } |
| } |
| if (builtin_type == "zeros" || builtin_type == "ones") { |
| if ((values.size() != 1) && !has_shape) { |
| throw ErrorReport(apply) |
| << "Built-in " << builtin_type |
| << " requires either 1 input or 1 shape attribute"; |
| } |
| } else { |
| // zeros_like or ones_like |
| if (values.size() != 1) { |
| throw ErrorReport(apply) |
| << "Built-in " << builtin_type << " requires 1 input"; |
| } |
| } |
| |
| auto op = cur().add_op(); |
| op->set_type("ConstantFill"); |
| if (values.size()) { |
| op->add_input(values[0]); |
| auto* input_as_shape = op->add_arg(); |
| input_as_shape->set_name("input_as_shape"); |
| if (builtin_type.find("_like") != std::string::npos) { |
| // zeros_like, ones_like take the shape of the input as constant |
| // tensor shape |
| input_as_shape->set_i(0); |
| } else { |
| // zeros, ones take the values in the tensor as constant tensor |
| // shape |
| input_as_shape->set_i(1); |
| } |
| } else { |
| fillArg(op->add_arg(), apply.attributes()[0]); |
| } |
| |
| auto value = op->add_arg(); |
| value->set_name("value"); |
| if (builtin_type.find("ones") != std::string::npos) { |
| value->set_f(1.0f); |
| } else { |
| value->set_f(0.0f); |
| } |
| appendOutputs(apply, op, outputs, 1); |
| } |
| // emitModule doesn't actually do anything except for allow |
| // statements like a = Module() to register 'a' as a valid identifier |
| // so that a.b = ... will work |
| void emitModule(const Apply& apply, const std::vector<std::string>& outputs) { |
| expectOutputs(apply, outputs, 1); |
| } |
| std::unordered_map< |
| std::string, |
| std::function<void( |
| DefCompiler*, |
| const Apply&, |
| const std::vector<std::string>& outputs)>> |
| builtins{{"zeros", &DefCompiler::emitFillOp}, |
| {"zeros_like", &DefCompiler::emitFillOp}, |
| {"ones", &DefCompiler::emitFillOp}, |
| {"ones_like", &DefCompiler::emitFillOp}, |
| {"Module", &DefCompiler::emitModule}}; |
| }; |
| |
| struct CompilationUnitImpl { |
| void defineFunction(const Def& def) { |
| if (functions.count(def.name().name()) > 0) { |
| throw ErrorReport(def) << def.name().name() << " already defined."; |
| } |
| DefCompiler c( |
| functions.emplace(def.name().name(), FunctionDefinition(def)) |
| .first->second, |
| functions); |
| c.run(); |
| } |
| |
| void define(const std::string& str) { |
| Parser p(str); |
| while (p.lexer().cur().kind != TK_EOF) { |
| defineFunction(Def(p.parseFunction())); |
| } |
| } |
| |
| std::unique_ptr<NetBase> createNet(Workspace* ws, const std::string& str) { |
| if (functions.count(str) == 0) |
| throw ErrorReport() << "undefined function: " << str << "\n"; |
| auto& def = functions.at(str); |
| return caffe2::CreateNet(*def.net_def, ws); |
| } |
| |
| void defineExtern(const std::string& name, std::unique_ptr<NetDef> net_def) { |
| // TODO: unify extern and function namespaces |
| if (functions.count(name) > 0) { |
| throw ErrorReport() << "function '" << name << "' already defined."; |
| } |
| functions.emplace(name, FunctionDefinition(std::move(net_def))); |
| } |
| |
| std::string getProto(const std::string& functionName) { |
| return functions.at(functionName).net_def->DebugString(); |
| } |
| |
| private: |
| friend struct DefCompiler; |
| SymbolTable functions; |
| }; |
| |
| CompilationUnit::CompilationUnit() : pImpl(new CompilationUnitImpl()) {} |
| |
| void CompilationUnit::define(const std::string& str) { |
| return pImpl->define(str); |
| } |
| |
| void CompilationUnit::defineExtern( |
| const std::string& name, |
| std::unique_ptr<NetDef> nd) { |
| pImpl->defineExtern(name, std::move(nd)); |
| } |
| |
| std::unique_ptr<NetBase> CompilationUnit::createNet( |
| Workspace* ws, |
| const std::string& str) { |
| return pImpl->createNet(ws, str); |
| } |
| |
| std::string CompilationUnit::getProto(const std::string& functionName) const { |
| return pImpl->getProto(functionName); |
| } |
| |
| CompilationUnit::~CompilationUnit() {} |
| |
| } // namespace script |
| } // namespace caffe2 |
