function.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef TVM_RELAY_FUNCTION_H_
#define TVM_RELAY_FUNCTION_H_
 
#include <tvm/ir/function.h>
#include <tvm/relay/expr.h>
 
#include <string>
 
namespace tvm {
namespace relay {
 
class FunctionNode : public BaseFuncNode {
 public:
  tvm::Array<Var> params;
  Expr body;
  Type ret_type;
  tvm::Array<TypeVar> type_params;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("params", &params);
    v->Visit("body", &body);
    v->Visit("ret_type", &ret_type);
    v->Visit("type_params", &type_params);
    v->Visit("attrs", &attrs);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }
 
  bool SEqualReduce(const FunctionNode* other, SEqualReducer equal) const {
    // Important to make def equal first.
    equal->MarkGraphNode();
    return equal.DefEqual(params, other->params) &&
           equal.DefEqual(type_params, other->type_params) && equal(ret_type, other->ret_type) &&
           equal(attrs, other->attrs) && equal(body, other->body);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce.DefHash(params);
    hash_reduce.DefHash(type_params);
    hash_reduce(ret_type);
    hash_reduce(attrs);
    hash_reduce(body);
  }
 
  TVM_DLL FuncType func_type_annotation() const;
 
  static constexpr const char* _type_key = "relay.Function";
  TVM_DECLARE_FINAL_OBJECT_INFO(FunctionNode, BaseFuncNode);
};
 
class Function : public BaseFunc {
 public:
  TVM_DLL Function(tvm::Array<Var> params, Expr body, Type ret_type, tvm::Array<TypeVar> ty_params,
                   tvm::DictAttrs attrs = DictAttrs(), Span span = Span());
 
  TVM_DEFINE_OBJECT_REF_METHODS(Function, BaseFunc, FunctionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
};
 
Function WithFields(Function function, Optional<Array<Var>> opt_params = Optional<Array<Var>>(),
                    Optional<Expr> opt_body = Optional<Expr>(),
                    Optional<Type> opt_ret_type = Optional<Type>(),
                    Optional<Array<TypeVar>> opt_ty_params = Optional<Array<TypeVar>>(),
                    Optional<DictAttrs> opt_attrs = Optional<DictAttrs>(),
                    Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                    Optional<Span> opt_span = Optional<Span>());
 
/*
 * \brief Returns the Relay FunctionNode represented by base_func if it should be optimized,
 * otherwise returns nullptr.
 *
 * This means returns nullptr:
 *  - For PrimFuncs, since not Relay Functions.
 *  - For Functions marked for external compilation (with "Compiler").
 *  - For Functions marked as already having an external definition (with "ExternalSymbol").
 *  - For Functions marked as not to be optimized (with "SkipOptimization").
 *
 * TODO(mbs): Audit all enumerations of IRModule::functions to use this or some family of such.
 */
const FunctionNode* AsOptimizableFunctionNode(const BaseFunc& base_func);
 
namespace attr {
 
constexpr const char* kPrimitive = "Primitive";
 
constexpr const char* kExtern = "Extern";
 
constexpr const char* kCompiler = "Compiler";
 
constexpr const char* kClosure = "Closure";
constexpr const char* kParams = "__params__";
constexpr const char* kSkipOptimization = "SkipOptimization";
constexpr const char* kComposite = "Composite";
constexpr const char* kInline = "Inline";
constexpr const char* kPartitionedFromPattern = "PartitionedFromPattern";
constexpr const char* kReshapeOnly = "relay.reshape_only";
 
}  // namespace attr
 
}  // namespace relay
}  // namespace tvm
#endif  // TVM_RELAY_FUNCTION_H_