<commit_before>/*
 * 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/graph_codegen.cc
 * \brief Graph runtime codegen
 */

#include <tvm/ir/module.h>
#include <tvm/relay/expr_functor.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/object.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/function.h>
#include <tvm/tir/stmt.h>

#include <algorithm>
#include <list>
#include <string>
#include <vector>

#include "compile_engine.h"
#include "utils.h"

namespace tvm {
namespace relay {
namespace backend {

using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;

class AotReturnSidVisitor : public ExprVisitor {
 public:
  explicit AotReturnSidVisitor(Map<Expr, Array<IntegerArray>> storage_device_map)
      : storage_device_map_{storage_device_map}, return_sid_{-1} {}

  IntegerArray FindReturnSid(Function func) {
    VisitExpr(func->body);
    return return_sid_;
  }

 protected:
  void AssignReturnSid(Expr e) {
    auto iter = storage_device_map_.find(e);
    if (iter != storage_device_map_.end()) {
      return_sid_ = (*iter).second[0];
    }
  }

  void VisitExpr_(const ConstantNode* cn) override {
    ExprVisitor::VisitExpr_(cn);
    AssignReturnSid(GetRef<Expr>(cn));
  }

  void VisitExpr_(const VarNode* vn) override {
    ExprVisitor::VisitExpr_(vn);
    AssignReturnSid(GetRef<Expr>(vn));
  }

  void VisitExpr_(const CallNode* cn) override {
    ExprVisitor::VisitExpr_(cn);
    AssignReturnSid(GetRef<Expr>(cn));
  }

  void VisitExpr_(const LetNode* op) override { VisitExpr(op->body); }

  void VisitExpr_(const TupleNode* tn) override {
    ExprVisitor::VisitExpr_(tn);
    AssignReturnSid(GetRef<Expr>(tn));
  }

 private:
  Map<Expr, Array<IntegerArray>> storage_device_map_;
  IntegerArray return_sid_;
};

/*! \brief Code generator for AOT executor */
class AOTExecutorCodegen : public ExprVisitor {
 protected:
  /*!
   * \brief Utility function to allocate a DLTensor or TVMValue
   * \param  type the type of allocation
   * \param num the number of variable to allocate on the stack
   * \return PrimExpr representing the allocated object
   */
  PrimExpr StackAlloca(std::string type, size_t num) {
    Array<PrimExpr> args = {tir::StringImm(type), ConstInt32(num)};
    return tir::Call(DataType::Handle(), tir::builtin::tvm_stack_alloca(), args);
  }

  /*!
   * \brief Utility function to convert a concrete integer to a PrimExpr.
   * \param num the number to convert
   * \return PrimExpr representing num
   */
  inline PrimExpr ConstInt32(size_t num) {
    ICHECK_LE(num, std::numeric_limits<int>::max());
    return tir::make_const(DataType::Int(32), static_cast<int>(num));
  }

  /*!
   * \brief Return a vector of variables that represents the sids for the given Relay Expr
   */
  std::vector<tir::Var> PackSid(Expr expr) {
    Array<IntegerArray> sids = storage_device_map_[expr];
    std::vector<tir::Var> sid_vars;

    // Note that an expression can have multiple sids associated with it
    // e.g., returning multiple values from a function
    for (const auto& sid : sids[0]) {
      // Determine if an sid is an output buffer
      int sid_int = static_cast<int>((sid.as<IntImmNode>())->value);
      auto output_iter = std::find(return_sid_.begin(), return_sid_.end(), sid_int);
      if (output_iter != return_sid_.end()) {
        int output_index = std::distance(return_sid_.begin(), output_iter);
        sid_vars.push_back(main_signature_[input_vars_.size() + output_index]);
        continue;
      }
      // Pack the sid inside the TVMValue
      auto sid_array = te::Var(MakeString("sid_", sid, "_value"), DataType::Handle());
      auto sid_value = sids_table_[sid];
      tvm::PrimExpr set_tensor =
          tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
                         {sid_array, 0, tir::builtin::kArrData, sid_value});
      stmts_.push_back(tir::LetStmt(sid_array, StackAlloca("array", 1), tir::Evaluate(set_tensor)));
      sid_vars.push_back(sid_array);
    }
    return sid_vars;
  }

  /*!
   * \brief Utility function to return a parameter associated with an expression
   * \param expr Relay Expression assicated with the parameter
   * \return Variable that represents the DLTensor associated with the parameters
   */
  tir::Var PackParam(Expr expr) {
    int param_sid = param_storage_ids_[params_by_expr_[expr]];
    auto param_array = te::Var(MakeString("param_", param_sid, "_array"), DataType::Handle());

    // Compose the lookup_call using a local stack
    Array<tir::Stmt> lookup_call;
    // Set the param to the value returned by lookup_call
    auto param_handle = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::lookup_param(),
                                       {tir::StringImm(params_by_expr_[expr])});

    tvm::PrimExpr set_param_array =
        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
                       {param_array, 0, tir::builtin::kArrData, param_handle});
    lookup_call.push_back(tir::Evaluate(set_param_array));

    tir::Stmt lookup_body = tir::SeqStmt(lookup_call);
<commit_msg>Fix typo in a comment (#8129)<commit_after>/*
 * 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/graph_codegen.cc
 * \brief Graph runtime codegen
 */

#include <tvm/ir/module.h>
#include <tvm/relay/expr_functor.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/object.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/function.h>
#include <tvm/tir/stmt.h>

#include <algorithm>
#include <list>
#include <string>
#include <vector>

#include "compile_engine.h"
#include "utils.h"

namespace tvm {
namespace relay {
namespace backend {

using IntegerArray = Array<Integer>;
using TargetsMap = std::unordered_map<int, Target>;

class AotReturnSidVisitor : public ExprVisitor {
 public:
  explicit AotReturnSidVisitor(Map<Expr, Array<IntegerArray>> storage_device_map)
      : storage_device_map_{storage_device_map}, return_sid_{-1} {}

  IntegerArray FindReturnSid(Function func) {
    VisitExpr(func->body);
    return return_sid_;
  }

 protected:
  void AssignReturnSid(Expr e) {
    auto iter = storage_device_map_.find(e);
    if (iter != storage_device_map_.end()) {
      return_sid_ = (*iter).second[0];
    }
  }

  void VisitExpr_(const ConstantNode* cn) override {
    ExprVisitor::VisitExpr_(cn);
    AssignReturnSid(GetRef<Expr>(cn));
  }

  void VisitExpr_(const VarNode* vn) override {
    ExprVisitor::VisitExpr_(vn);
    AssignReturnSid(GetRef<Expr>(vn));
  }

  void VisitExpr_(const CallNode* cn) override {
    ExprVisitor::VisitExpr_(cn);
    AssignReturnSid(GetRef<Expr>(cn));
  }

  void VisitExpr_(const LetNode* op) override { VisitExpr(op->body); }

  void VisitExpr_(const TupleNode* tn) override {
    ExprVisitor::VisitExpr_(tn);
    AssignReturnSid(GetRef<Expr>(tn));
  }

 private:
  Map<Expr, Array<IntegerArray>> storage_device_map_;
  IntegerArray return_sid_;
};

/*! \brief Code generator for AOT executor */
class AOTExecutorCodegen : public ExprVisitor {
 protected:
  /*!
   * \brief Utility function to allocate a DLTensor or TVMValue
   * \param  type the type of allocation
   * \param num the number of variable to allocate on the stack
   * \return PrimExpr representing the allocated object
   */
  PrimExpr StackAlloca(std::string type, size_t num) {
    Array<PrimExpr> args = {tir::StringImm(type), ConstInt32(num)};
    return tir::Call(DataType::Handle(), tir::builtin::tvm_stack_alloca(), args);
  }

  /*!
   * \brief Utility function to convert a concrete integer to a PrimExpr.
   * \param num the number to convert
   * \return PrimExpr representing num
   */
  inline PrimExpr ConstInt32(size_t num) {
    ICHECK_LE(num, std::numeric_limits<int>::max());
    return tir::make_const(DataType::Int(32), static_cast<int>(num));
  }

  /*!
   * \brief Return a vector of variables that represents the sids for the given Relay Expr
   */
  std::vector<tir::Var> PackSid(Expr expr) {
    Array<IntegerArray> sids = storage_device_map_[expr];
    std::vector<tir::Var> sid_vars;

    // Note that an expression can have multiple sids associated with it
    // e.g., returning multiple values from a function
    for (const auto& sid : sids[0]) {
      // Determine if an sid is an output buffer
      int sid_int = static_cast<int>((sid.as<IntImmNode>())->value);
      auto output_iter = std::find(return_sid_.begin(), return_sid_.end(), sid_int);
      if (output_iter != return_sid_.end()) {
        int output_index = std::distance(return_sid_.begin(), output_iter);
        sid_vars.push_back(main_signature_[input_vars_.size() + output_index]);
        continue;
      }
      // Pack the sid inside the TVMValue
      auto sid_array = te::Var(MakeString("sid_", sid, "_value"), DataType::Handle());
      auto sid_value = sids_table_[sid];
      tvm::PrimExpr set_tensor =
          tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
                         {sid_array, 0, tir::builtin::kArrData, sid_value});
      stmts_.push_back(tir::LetStmt(sid_array, StackAlloca("array", 1), tir::Evaluate(set_tensor)));
      sid_vars.push_back(sid_array);
    }
    return sid_vars;
  }

  /*!
   * \brief Utility function to return a parameter associated with an expression
   * \param expr Relay Expression associated with the parameter
   * \return Variable that represents the DLTensor associated with the parameters
   */
  tir::Var PackParam(Expr expr) {
    int param_sid = param_storage_ids_[params_by_expr_[expr]];
    auto param_array = te::Var(MakeString("param_", param_sid, "_array"), DataType::Handle());

    // Compose the lookup_call using a local stack
    Array<tir::Stmt> lookup_call;
    // Set the param to the value returned by lookup_call
    auto param_handle = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::lookup_param(),
                                       {tir::StringImm(params_by_expr_[expr])});

    tvm::PrimExpr set_param_array =
        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
                       {param_array, 0, tir::builtin::kArrData, param_handle});
    lookup_call.push_back(tir::Evaluate(set_param_array));

    tir::Stmt lookup_body = tir::SeqStmt(lookup_call);

{'fraction_non_alphanumeric': 0.11993285774234158, 'fraction_numerical': 0.0030214015946286193, 'mean_word_length': 3.956738768718802, 'pattern_counts': {'":': 0, '<': 67, '<?xml version=': 0, '>': 73, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4030304', 'n_tokens_mistral': 3744, 'n_tokens_neox': 3685, 'n_words': 995}