src/tirx/analysis/stmt_finding.cc

/*
 * 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.
 */
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt_functor.h>

namespace tvm {
namespace tirx {

const PrimFuncNode* FindEntryFunc(const IRModule& mod, GlobalVar* result_g_var) {
  GlobalVar result;
  // Priority 1: PrimFunc marked as `tirx::attr::kIsEntryFunc`
  int num_prim_func = 0;
  const tirx::PrimFuncNode* main_func = nullptr;
  const tirx::PrimFuncNode* last_func = nullptr;
  for (const auto& kv : mod->functions) {
    GlobalVar gv = kv.first;
    BaseFunc base_func = kv.second;
    if (const auto* func = base_func.as<tirx::PrimFuncNode>()) {
      last_func = func;
      if (func->HasNonzeroAttr(tirx::attr::kIsEntryFunc)) {
        if (result_g_var != nullptr) {
          *result_g_var = gv;
        }
        return func;
      }
      if (gv->name_hint == "main") {
        main_func = func;
        result = gv;
      }
      ++num_prim_func;
    }
  }
  // Priority 2: PrimFunc whose name is `main`
  if (main_func != nullptr) {
    if (result_g_var != nullptr) {
      *result_g_var = result;
    }
    return main_func;
  }
  // Priority 3: The only PrimFunc in the IRModule
  if (num_prim_func == 1) {
    if (result_g_var != nullptr) {
      *result_g_var = result;
    }
    return last_func;
  }
  return nullptr;
}

}  // namespace tirx
}  // namespace tvm