zhiics commented on a change in pull request #5045: [REFACTOR] Streamline Function Attr interface. URL: https://github.com/apache/incubator-tvm/pull/5045#discussion_r391342446
########## File path: include/tvm/ir/function.h ########## @@ -0,0 +1,119 @@ +/* + * 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. + */ + +/*! + * \file tvm/ir/expr.h + * \brief Base expr nodes in TVM. + */ +#ifndef TVM_IR_FUNCTION_H_ +#define TVM_IR_FUNCTION_H_ + +#include <tvm/ir/expr.h> +#include <tvm/ir/attrs.h> +#include <type_traits> +#include <string> + + +namespace tvm { + +/*! + * \brief Base node of all functions. + * + * We support several variants of functions throughout the stack. + * All of the functions shares the same type system(via checked_type) + * to support cross variant calls. + * + * \sa BaseFunc + */ +class BaseFuncNode : public RelayExprNode { + public: + /*! \brief Additional attributes storing the meta-data */ + DictAttrs attrs; + + /*! + * \brief Get a function attribute. + * + * \param attr_key The attribute key. + * \param default_value The default value if the key does not exist, defaults to nullptr. + * + * \return The result + * + * \tparam TOBjectRef the expected object type. + * \throw Error if the key exists but the value does not match TObjectRef + * + * \code + * + * void GetAttrExample(const BaseFunc& f) { + * Integer value = f->GetAttr<Integer>("AttrKey", 0); + * } + * + * \endcode + */ + template<typename TObjectRef> + TObjectRef GetAttr(const std::string& attr_key, + TObjectRef default_value = NullValue<TObjectRef>()) const { + static_assert(std::is_base_of<ObjectRef, TObjectRef>::value, + "Can only call GetAttr with ObjectRef types."); + if (!attrs.defined()) return default_value; + auto it = attrs->dict.find(attr_key); + if (it != attrs->dict.end()) { + return Downcast<TObjectRef>((*it).second); + } else { + return default_value; + } + } + + /*! + * \brief Check whether the function has an non-zero integer attr. + * + * This function can be used to check whether an optional + * attribute mark(e.g. inline) exists. + * + * \param attr_key The key to the attribute. + * \return The check result. + * + * \code + * + * void HasNonzeroAttrExample(const BaseFunc& f) { + * if (f->HasNonzeroAttr(attr::Inline)) { Review comment: kInline ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
