areusch commented on code in PR #12087: URL: https://github.com/apache/tvm/pull/12087#discussion_r931325454
########## python/tvm/relay/backend/contrib/uma/_template/passes.py: ########## @@ -0,0 +1,137 @@ +# 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. +"""Transform passes for the my_ai_hw accelerator""" + +import tvm +from tvm import relay, tir +from tvm.relay.backend.contrib.uma.api.utils import add_llvm_to_block + + [email protected]_func_pass(opt_level=2) +class MyAiHwConv2dPass: + def transform_function( + self, func: tvm.tir.PrimFunc, mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext + ) -> tvm.tir.PrimFunc: + return self._my_ai_hw_conv2d_pass(func, mod, ctx) + + @staticmethod + def _my_ai_hw_conv2d_pass(func, mod, ctx): + _found_blocks = [] + _loops = dict() + _handles = [] + _entry_node = None + _external_function_name = "my_ai_hw_conv2dnchw" Review Comment: `@classmethod` should fix this ########## src/relay/backend/contrib/uma/relay_to_tir.cc: ########## @@ -0,0 +1,174 @@ +/* + * 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/uma/codegen.cc + * + * \brief this file contains the target hooks for the Universal Modular Accelerator Interface (UMA). + */ + +#include <tvm/ir/error.h> +#include <tvm/relay/analysis.h> +#include <tvm/relay/attrs/annotation.h> +#include <tvm/relay/expr.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/transform.h> +#include <tvm/target/target.h> +#include <tvm/tir/function.h> + +#include <unordered_map> +#include <unordered_set> +#include <utility> +#include <vector> + +namespace tvm { +namespace relay { +namespace contrib { +namespace uma { + +/*! + * \brief This mutator outlines functions that are marked with a named + * "Compiler" attribute. Functions that do not match this condition remain + * unaltered. + */ +class OutlineCompilerFunctionsMutator : public MixedModeMutator { + public: + explicit OutlineCompilerFunctionsMutator(const IRModule& mod, const std::string& compiler_name) + : mod_(mod), compiler_name_(compiler_name) {} + + Expr VisitExpr_(const LetNode* op) final { + auto pre_visit = [this](const LetNode* op) { + Expr var = this->VisitExpr(op->var); + Expr value = this->VisitExpr(op->value); + + // Outlineable function no longer needs let binding + if (this->CanOutlineExpr(value)) { + this->memo_[var] = value; + } + }; + auto post_visit = [this](const LetNode* op) { + // Rely on the Memoizer to cache pre-visit values + Expr value = this->VisitExpr(op->value); + Expr body = this->VisitExpr(op->body); + auto expr = GetRef<Expr>(op); + + // Drop the let binding + if (this->CanOutlineExpr(value)) { + this->memo_[expr] = this->VisitExpr(op->body); + } else { + Var var = Downcast<Var>(this->VisitExpr(op->var)); + if (var.same_as(op->var) && value.same_as(op->value) && body.same_as(op->body)) { + this->memo_[expr] = expr; + } else { + this->memo_[expr] = Let(var, value, body); + } + } + }; + ExpandANormalForm(op, pre_visit, post_visit); + return memo_[GetRef<Expr>(op)]; + } + + Expr Rewrite_(const CallNode* pre, const Expr& post) override { + Call call = Downcast<Call>(post); + if (CanOutlineExpr(call->op)) { + Function func = Downcast<Function>(call->op); + auto gv_name = func->GetAttr<String>("global_symbol").value_or(""); + ICHECK_NE(gv_name, "") + << "Function to be outlined must have global_symbol attribute, but didn't."; + GlobalVar gv(gv_name); + if (func->checked_type_.defined()) { + gv->checked_type_ = func->checked_type(); + } + mod_->Update(gv, func); + return Call(gv, call->args, call->attrs, call->type_args); + } + return post; + } + + private: + /*! + * \brief Check if the expr is a function and has the same + * compiler name as compiler_name_. + * + * \param expr The input expr. + * \return True if is outlineable else False. + */ + bool CanOutlineExpr(const Expr& expr) { + if (!expr->IsInstance<FunctionNode>()) { + return false; + } + Function func = Downcast<Function>(expr); + auto compiler = func->GetAttr<String>(attr::kCompiler); + if (!compiler.defined()) { + return false; + } + if (compiler != compiler_name_) { + return false; + } + return true; + } + + /*! \brief The module that the pass will run on. */ + IRModule mod_; + /*! \brief The name of the compiler to enable outlining on external functions for. */ + std::string compiler_name_; +}; + +/*! + * \brief A pass to outline compiler specific functions. + */ +tvm::transform::Pass OutlineCompilerFunctions(const std::string& compiler_name) { + runtime::TypedPackedFunc<IRModule(IRModule, transform::PassContext)> pass_func = + [=](IRModule mod, transform::PassContext ctx) { + GlobalVar gv = mod->GetGlobalVar("main"); + Function main_func = Downcast<Function>(mod->Lookup("main")); + auto new_main_body = + OutlineCompilerFunctionsMutator(mod, compiler_name).VisitExpr(main_func->body); + if (!new_main_body.same_as(main_func->body)) { + Function new_main_func = WithFields(main_func, main_func->params, new_main_body); + mod->Update(gv, new_main_func); + } + return mod; + }; + return tvm::transform::CreateModulePass(pass_func, 0, + "relay.backend.contrib.uma.OutlineCompilerFunctions", {}); +} + +TVM_REGISTER_GLOBAL("relay.ext.uma.OutlineCompilerFunctions") + .set_body_typed(OutlineCompilerFunctions); + +/*! + * \brief This pass will lower NPU functions in a Relay module to scheduled TIR prim functions. Review Comment: NPU ########## src/relay/backend/contrib/uma/relay_to_tir.cc: ########## @@ -0,0 +1,174 @@ +/* + * 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/uma/codegen.cc + * + * \brief this file contains the target hooks for the Universal Modular Accelerator Interface (UMA). + */ + +#include <tvm/ir/error.h> +#include <tvm/relay/analysis.h> +#include <tvm/relay/attrs/annotation.h> +#include <tvm/relay/expr.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/transform.h> +#include <tvm/target/target.h> +#include <tvm/tir/function.h> + +#include <unordered_map> +#include <unordered_set> +#include <utility> +#include <vector> + +namespace tvm { +namespace relay { +namespace contrib { +namespace uma { + +/*! + * \brief This mutator outlines functions that are marked with a named + * "Compiler" attribute. Functions that do not match this condition remain + * unaltered. + */ +class OutlineCompilerFunctionsMutator : public MixedModeMutator { Review Comment: since we're duplicating it here, i'd prefer if we could factor it out now ########## tests/python/contrib/test_uma/test_partition.py: ########## @@ -0,0 +1,71 @@ +# 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. + +import pytest + +import tvm + +from tvm.relay.backend.contrib.uma.api import UMAPartitioner +from tvm.relay.op.contrib.register import get_pattern_table +from tvm.relay.testing import resnet, mlp + + +def test_partition_table(): + partitioner = UMAPartitioner("test_partition") + assert get_pattern_table("test_partition") is None + + partitioner.register() + + assert get_pattern_table("test_partition") is not None + + [email protected]( + "workload,backend,merge,expected_partitions", + [ + ("resnet", "dnnl", False, 17), + ("resnet", "dnnl", True, 17), Review Comment: ah could we either document in a comment where the number 17 came from or derive it in code (can still assert against the number 17, just want to help future test-editors know how to keep this test updated). ########## python/tvm/relay/backend/contrib/uma/_template/backend.py: ########## @@ -0,0 +1,53 @@ +# 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. +"""UMA backend for the my_ai_hw accelerator""" +from .passes import MyAiHwConv2dPass +from ..api.utils import PassPhase +from ..backend import UMABackend +from .codegen import gen_includes, gen_replace_call_extern +from .patterns import conv2d_pattern + + +class MyAiHwBackend(UMABackend): + """UMA backend for the MyAiHw accelerator.""" + + def __init__(self): + super().__init__() + + ####################################################################### + # Target configuration + ####################################################################### + self._register_target_attr("dimension") + + ####################################################################### + # Relay to Relay function registration + ####################################################################### + self._register_pattern("conv2d", conv2d_pattern()) + + ####################################################################### + # Relay to TIR function registration + ####################################################################### + self._register_tir_pass(PassPhase.TIR_PHASE_0, MyAiHwConv2dPass()) Review Comment: cc @mbaret one thing we've been discussing is that it might make sense for these enums to roughly correspond to "guarantees about the form of the IRModule" (e.g. "all the primitive Functions are lowered" or "Buffers are defined"). To that end I think it would be great to avoid referencing e.g. TIR_PHASE_n in the enum unless those are distinct guarantees. -- 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]
