mbs-octoml commented on a change in pull request #8806: URL: https://github.com/apache/tvm/pull/8806#discussion_r706249139
########## File path: python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py ########## @@ -0,0 +1,277 @@ +# 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. +# pylint: disable=invalid-name, unused-argument +"""Different schedulers for Arm(R) Ethos(TM)-U NPU""" Review comment: I'm starting to see the forest for the trees but a few paragraphs explaining the scheduling strategy would really help (along with a pointer to the RFC). The main points I think are: - make the dma with its layout xforms explicit (which in an ideal world could be done without having to replicate the inner conv2d TE defns?) - hoist constants - prepare for cascading Is that about right? I'm not sure what we should say about how this plays with the new 'schedulable' TIR. We are prepared to port at some point? TVM overall is on the hook for preserving the 'legacy' scheduling API? Perhaps a GitHub issue is in order. ########## File path: src/relay/backend/contrib/ethosu/to_te_graph.cc ########## @@ -0,0 +1,234 @@ +/* + * 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 relay/backend/contrib/ethosu/to_te_graph.cc + * \brief Lower a Relay function to a TE graph. Review comment: Can I check this file (and it's dependence on lower_call in compile_engine.py) is temporary scaffolding pending a suitable hook mechanism? ########## File path: src/relay/backend/contrib/ethosu/to_te_graph.cc ########## @@ -0,0 +1,234 @@ +/* + * 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 relay/backend/contrib/ethosu/to_te_graph.cc + * \brief Lower a Relay function to a TE graph. + */ +#include <tvm/driver/driver_api.h> +#include <tvm/ir/type_functor.h> +#include <tvm/relay/analysis.h> +#include <tvm/relay/attrs/device_copy.h> +#include <tvm/relay/expr.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/op.h> +#include <tvm/relay/op_attr_types.h> +#include <tvm/runtime/registry.h> +#include <tvm/te/operation.h> +#include <tvm/te/schedule.h> +#include <tvm/te/schedule_pass.h> +#include <tvm/topi/tags.h> + +#include <functional> +#include <limits> +#include <mutex> +#include <unordered_map> +#include <utility> +#include <vector> + +#include "../../compile_engine.h" +#include "../../utils.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace ethosu { + +/*! \brief Node container to represent a Tensor Expression graph. */ +class TEGraphNode : public Object { + public: + /* \brief The inputs to the graph */ + tvm::Array<te::Tensor> inputs; + /* \brief The outputs to the graph */ + tvm::Array<te::Tensor> outputs; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("inputs", &inputs); + v->Visit("outputs", &outputs); + } + + static constexpr const char* _type_key = "relay.TEGraph"; + TVM_DECLARE_FINAL_OBJECT_INFO(TEGraphNode, Object); +}; + +class TEGraph : public ObjectRef { + public: + TVM_DEFINE_OBJECT_REF_METHODS(TEGraph, ObjectRef, TEGraphNode); +}; + +TVM_REGISTER_NODE_TYPE(TEGraphNode); + +Array<IndexExpr> GetShape(const Array<IndexExpr>& shape) { + // for now, we always use int32 shape when possible + // even if the result of shape inference becomes int64. + Array<IndexExpr> res; + for (IndexExpr val : shape) { + const int64_t* pval = tir::as_const_int(val); + if (pval != nullptr) { +#ifndef TVM_INDEX_DEFAULT_I64 + ICHECK_LE(pval[0], std::numeric_limits<int32_t>::max()); + ICHECK_GE(pval[0], std::numeric_limits<int32_t>::min()); + res.push_back(IntImm(DataType::Int(32), *pval)); +#else + res.push_back(val); +#endif // TVM_INDEX_DEFAULT_I64 + } else if (val->IsInstance<tir::AnyNode>()) { + res.push_back(val.as<tir::AnyNode>()->ToVar()); + } else { + res.push_back(val); + } + } + return res; +} + +class RelayToTE : public backend::MemoizedExprTranslator<Array<te::Tensor>> { + public: + RelayToTE() = default; + + TEGraph Lower(const Function& prim_func) { + auto graph_node = make_object<TEGraphNode>(); + for (Var param : prim_func->params) { + Array<tvm::te::Tensor> inputs; + if (const auto* ttype = param->checked_type().as<TensorTypeNode>()) { Review comment: You can first relay::FlattenTupleType then just itr over the resut. ########## File path: src/relay/backend/contrib/ethosu/to_te_graph.cc ########## @@ -0,0 +1,234 @@ +/* + * 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 relay/backend/contrib/ethosu/to_te_graph.cc + * \brief Lower a Relay function to a TE graph. + */ +#include <tvm/driver/driver_api.h> +#include <tvm/ir/type_functor.h> +#include <tvm/relay/analysis.h> +#include <tvm/relay/attrs/device_copy.h> +#include <tvm/relay/expr.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/op.h> +#include <tvm/relay/op_attr_types.h> +#include <tvm/runtime/registry.h> +#include <tvm/te/operation.h> +#include <tvm/te/schedule.h> +#include <tvm/te/schedule_pass.h> +#include <tvm/topi/tags.h> + +#include <functional> +#include <limits> +#include <mutex> +#include <unordered_map> +#include <utility> +#include <vector> + +#include "../../compile_engine.h" +#include "../../utils.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace ethosu { + +/*! \brief Node container to represent a Tensor Expression graph. */ +class TEGraphNode : public Object { + public: + /* \brief The inputs to the graph */ + tvm::Array<te::Tensor> inputs; + /* \brief The outputs to the graph */ + tvm::Array<te::Tensor> outputs; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("inputs", &inputs); + v->Visit("outputs", &outputs); + } + + static constexpr const char* _type_key = "relay.TEGraph"; + TVM_DECLARE_FINAL_OBJECT_INFO(TEGraphNode, Object); +}; + +class TEGraph : public ObjectRef { + public: + TVM_DEFINE_OBJECT_REF_METHODS(TEGraph, ObjectRef, TEGraphNode); +}; + +TVM_REGISTER_NODE_TYPE(TEGraphNode); + +Array<IndexExpr> GetShape(const Array<IndexExpr>& shape) { + // for now, we always use int32 shape when possible + // even if the result of shape inference becomes int64. + Array<IndexExpr> res; + for (IndexExpr val : shape) { + const int64_t* pval = tir::as_const_int(val); + if (pval != nullptr) { +#ifndef TVM_INDEX_DEFAULT_I64 + ICHECK_LE(pval[0], std::numeric_limits<int32_t>::max()); + ICHECK_GE(pval[0], std::numeric_limits<int32_t>::min()); + res.push_back(IntImm(DataType::Int(32), *pval)); +#else + res.push_back(val); +#endif // TVM_INDEX_DEFAULT_I64 + } else if (val->IsInstance<tir::AnyNode>()) { + res.push_back(val.as<tir::AnyNode>()->ToVar()); + } else { + res.push_back(val); + } + } + return res; +} + +class RelayToTE : public backend::MemoizedExprTranslator<Array<te::Tensor>> { + public: + RelayToTE() = default; + + TEGraph Lower(const Function& prim_func) { + auto graph_node = make_object<TEGraphNode>(); + for (Var param : prim_func->params) { + Array<tvm::te::Tensor> inputs; + if (const auto* ttype = param->checked_type().as<TensorTypeNode>()) { + tvm::te::Tensor tensor = tvm::te::placeholder(GetShape(ttype->shape), ttype->dtype); + graph_node->inputs.push_back(tensor); + inputs.push_back(tensor); + } else { + // flatten tuple of tensor type. + const auto* tuple_type = param->type_as<TupleTypeNode>(); + for (Type field : tuple_type->fields) { + const auto* ttype = field.as<TensorTypeNode>(); + ICHECK(ttype != nullptr); + tvm::te::Tensor tensor = tvm::te::placeholder(GetShape(ttype->shape), ttype->dtype); + graph_node->inputs.push_back(tensor); + inputs.push_back(tensor); + } + } + memo_[param] = inputs; + } + graph_node->outputs = this->VisitExpr(prim_func->body); + return TEGraph(graph_node); + } + + Array<te::Tensor> VisitExpr_(const VarNode* op) final { + LOG(FATAL) << "Free variable " << op->name_hint(); + return {}; + } + + Array<te::Tensor> VisitExpr_(const ConstantNode* op) final { + using tir::make_const; + ICHECK(op->is_scalar()); + void* data = op->data->data; + DataType dtype = DataType(op->data->dtype); + auto value = te::compute( + {}, + [&](const Array<tvm::tir::Var>&) { + if (dtype == DataType::Int(32)) { + return make_const(dtype, static_cast<const int32_t*>(data)[0]); + } else if (dtype == DataType::Int(64)) { + return make_const(dtype, static_cast<const int64_t*>(data)[0]); + } else if (dtype == DataType::Float(32)) { + return make_const(dtype, static_cast<const float*>(data)[0]); + } else if (dtype == DataType::Float(64)) { + return make_const(dtype, static_cast<const double*>(data)[0]); + } else if (dtype == DataType::Bool()) { + return make_const(dtype, static_cast<const uint8_t*>(data)[0]); + } else { + LOG(FATAL) << "not handled"; + return tvm::PrimExpr(); + } + }, + "compile_engine_const", topi::kBroadcast); + return {value}; + } + + Array<te::Tensor> VisitExpr_(const CallNode* call_node) final { + static auto flower_call = tvm::runtime::Registry::Get("relay.backend.lower_call"); Review comment: Please comment that this API is pending removal in #8775. ########## File path: src/relay/backend/contrib/ethosu/to_te_graph.cc ########## @@ -0,0 +1,234 @@ +/* + * 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 relay/backend/contrib/ethosu/to_te_graph.cc + * \brief Lower a Relay function to a TE graph. Review comment: Looks like this PR has overtaken the relay_to_tir hook mechanism train. - Let's make abundantly clear in the comment this is a copy of relay::tec::ScheduleBuilder for the purpose of intercepting the scheduling and this is a temporary workaround. It would be a shame if a new contributor assumed this is the pattern for controlling scheduling. - And let's file a GitHub issue or point to the existing one for the rely_to_tir hook mechanism. -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
