<commit_before>#include "codegen_cpu.h"
#include "program.h"
#include <taichi/common/bit.h>

TLANG_NAMESPACE_BEGIN

void CPUCodeGen::visit_intrinsics(Expr &expr) {
  /*
  TC_P(expr->id);
  TC_P(expr->node_type_name());
  TC_P(num_groups);
  TC_P(expr->lanes);
  */
  auto vv_width = expr->lanes;
  TC_ASSERT(vv_width == 1 || vv_width == simd_width);
  int split = vv_width / simd_width;
  auto vv_type = [&](DataType dt) {
    if (expr->lanes == 1) {
      return fmt::format("vvec<{}, {}, {}>", data_type_name(dt), 1, 1);
    } else {
      return fmt::format("vvec<{}, {}, {}>", data_type_name(dt), simd_width,
                         split);
    }
  };
  if (expr->type == NodeType::addr) {
    return;
  }
  if (expr->var_name == "") {
    // expr->var_name = create_variable();
    expr->var_name = fmt::format("var{}", expr->id);
    /*
    TC_INFO("{} {} {} -> {}", expr->id, expr->node_type_name(),
            expr->data_type_name(), expr->var_name);
            */
  } else
    return;  // visited

  for (auto &m : expr->members) {
    if (!m->name().empty()) {
      emit_code("// @ {}", m->name());
    }
  }

  auto address_elements = [&](std::string index, int start = 1) {
    std::string ret = "";
    for (int i = start; i < (int)expr->ch.size(); i++) {
      ret += fmt::format(", {}.element({})", expr->ch[i]->var_name, index);
    }
    return ret;
  };

  if (expr->type == NodeType::binary) {
    auto op = binary_type_name(expr->binary_type);
    emit_code("auto {} = {}({}, {});", expr->var_name, op,
              expr->ch[0]->var_name, expr->ch[1]->var_name);
    // emit_code("{}.print();", expr->var_name);
  } else if (expr->type == NodeType::land) {
    TC_ASSERT(expr[1]->type == NodeType::imm)
    TC_WARN("member imm");
    emit_code("auto {} = land({}, {});", expr->var_name, expr[0]->var_name,
              expr[1]->value<int>());
  } else if (expr->type == NodeType::shr) {
    TC_WARN("member imm");
    TC_ASSERT(expr[1]->type == NodeType::imm)
    emit_code("auto {} = shr({}, {});", expr->var_name, expr[0]->var_name,
              expr[1]->value<int>());
  } else if (expr->type == NodeType::shl) {
    TC_WARN("member imm");
    TC_ASSERT(expr[1]->type == NodeType::imm)
    emit_code("auto {} = shl({}, {});", expr->var_name, expr[0]->var_name,
              expr[1]->value<int>());
  } else if (expr->type == NodeType::cmp) {
    auto t = expr->value<CmpType>();
    if (t == CmpType::ne) {
      emit_code("auto {} = cmp_ne({}, {});", expr->var_name, expr[0]->var_name,
                expr[1]->var_name);
    } else if (t == CmpType::lt) {
      emit_code("auto {} = cmp_lt({}, {});", expr->var_name, expr[0]->var_name,
                expr[1]->var_name);
    } else {
      TC_NOT_IMPLEMENTED
    }
  } else if (expr->type == NodeType::select) {
    emit_code("auto {} = select({}, {}, {});", expr->var_name,
              expr[0]->var_name, expr[1]->var_name, expr[2]->var_name);
  } else if (expr->type == NodeType::floor) {
    emit_code("auto {} = floor({});", expr->var_name, expr[0]->var_name);
  } else if (expr->type == NodeType::cast) {
    if (expr->data_type == DataType::i32) {
      emit_code("auto {} = {}.cast<int32>();", expr->var_name,
                expr[0]->var_name);
    } else if (expr->data_type == DataType::f32) {
      emit_code("auto {} = {}.cast<float32>();", expr->var_name,
                expr[0]->var_name);
    } else {
      TC_NOT_IMPLEMENTED
    }
  } else if (expr->type == NodeType::vload) {
    emit_code("auto {} = {}::load(access_{}(context.buffers[0] {}));",
              expr->var_name, vv_type(expr->data_type),
              expr[0]->new_addresses(0)->node_type_name, address_elements("0"));
  } else if (expr->type == NodeType::vstore) {
    emit_code("{}.store(access_{}(context.buffers[0] {}));", expr[1]->var_name,
              expr[0]->new_addresses(0)->node_type_name,
              address_elements("0", 2));
  } else if (expr->type == NodeType::load) {
    emit_code("auto {} = {}::load({});", expr->var_name,
              vv_type(expr->data_type), expr[0]->var_name);
#if (0)
    bool regular = false;
    if (expr[0]->type == NodeType::pointer &&
        expr[0][1]->type == NodeType::index) {
      regular = true;
    } else {
    }
    if (regular) {
      /*
      emit_code("auto {} = load<{}, {}>({}_base, {}_offsets);", expr->var_name,
                vv_width, data_type_name(expr->data_type), expr[0]->var_name,
                expr[0]->var_name);
                */
      // TC_P(expr->members.size());

      std::vector<int> offsets;
      for (int i = 0; i + 1 < (int)expr->members.size(); i++) {
        TC_ASSERT(
            expr->members[i]->addr().same_type(expr->members[i + 1]->addr()));
      }
      for (int i = 0; i < (int)expr->members.size(); i++) {
        offsets.push_back(expr->members[i]->addr().offset());<commit_msg>fixed vload/vstore with jumping index<commit_after>#include "codegen_cpu.h"
#include "program.h"
#include <taichi/common/bit.h>

TLANG_NAMESPACE_BEGIN

void CPUCodeGen::visit_intrinsics(Expr &expr) {
  /*
  TC_P(expr->id);
  TC_P(expr->node_type_name());
  TC_P(num_groups);
  TC_P(expr->lanes);
  */
  auto vv_width = expr->lanes;
  TC_ASSERT(vv_width == 1 || vv_width == simd_width);
  int split = vv_width / simd_width;
  auto vv_type = [&](DataType dt) {
    if (expr->lanes == 1) {
      return fmt::format("vvec<{}, {}, {}>", data_type_name(dt), 1, 1);
    } else {
      return fmt::format("vvec<{}, {}, {}>", data_type_name(dt), simd_width,
                         split);
    }
  };
  if (expr->type == NodeType::addr) {
    return;
  }
  if (expr->var_name == "") {
    // expr->var_name = create_variable();
    expr->var_name = fmt::format("var{}", expr->id);
    /*
    TC_INFO("{} {} {} -> {}", expr->id, expr->node_type_name(),
            expr->data_type_name(), expr->var_name);
            */
  } else
    return;  // visited

  for (auto &m : expr->members) {
    if (!m->name().empty()) {
      emit_code("// @ {}", m->name());
    }
  }

  auto address_elements = [&](SNode *snode, std::string index, int start = 1) {
    std::vector<std::string> elems(max_num_indices, ", 0");
    for (int j = start; j < (int)expr->ch.size(); j++) {
      elems[snode->index_order[j - start]] =
          fmt::format(", {}.element({})", expr->ch[j]->var_name, index);
    }
    std::string total_elem = "";
    for (int j = 0; j < max_num_indices; j++) {
      total_elem += elems[j];
    }

    return total_elem;
  };

  if (expr->type == NodeType::binary) {
    auto op = binary_type_name(expr->binary_type);
    emit_code("auto {} = {}({}, {});", expr->var_name, op,
              expr->ch[0]->var_name, expr->ch[1]->var_name);
    // emit_code("{}.print();", expr->var_name);
  } else if (expr->type == NodeType::land) {
    TC_ASSERT(expr[1]->type == NodeType::imm)
    TC_WARN("member imm");
    emit_code("auto {} = land({}, {});", expr->var_name, expr[0]->var_name,
              expr[1]->value<int>());
  } else if (expr->type == NodeType::shr) {
    TC_WARN("member imm");
    TC_ASSERT(expr[1]->type == NodeType::imm)
    emit_code("auto {} = shr({}, {});", expr->var_name, expr[0]->var_name,
              expr[1]->value<int>());
  } else if (expr->type == NodeType::shl) {
    TC_WARN("member imm");
    TC_ASSERT(expr[1]->type == NodeType::imm)
    emit_code("auto {} = shl({}, {});", expr->var_name, expr[0]->var_name,
              expr[1]->value<int>());
  } else if (expr->type == NodeType::cmp) {
    auto t = expr->value<CmpType>();
    if (t == CmpType::ne) {
      emit_code("auto {} = cmp_ne({}, {});", expr->var_name, expr[0]->var_name,
                expr[1]->var_name);
    } else if (t == CmpType::lt) {
      emit_code("auto {} = cmp_lt({}, {});", expr->var_name, expr[0]->var_name,
                expr[1]->var_name);
    } else {
      TC_NOT_IMPLEMENTED
    }
  } else if (expr->type == NodeType::select) {
    emit_code("auto {} = select({}, {}, {});", expr->var_name,
              expr[0]->var_name, expr[1]->var_name, expr[2]->var_name);
  } else if (expr->type == NodeType::floor) {
    emit_code("auto {} = floor({});", expr->var_name, expr[0]->var_name);
  } else if (expr->type == NodeType::cast) {
    if (expr->data_type == DataType::i32) {
      emit_code("auto {} = {}.cast<int32>();", expr->var_name,
                expr[0]->var_name);
    } else if (expr->data_type == DataType::f32) {
      emit_code("auto {} = {}.cast<float32>();", expr->var_name,
                expr[0]->var_name);
    } else {
      TC_NOT_IMPLEMENTED
    }
  } else if (expr->type == NodeType::vload) {
    emit_code("auto {} = {}::load(access_{}(context.buffers[0] {}));",
              expr->var_name, vv_type(expr->data_type),
              expr[0]->new_addresses(0)->node_type_name,
              address_elements(expr[0]->new_addresses(0), "0"));
  } else if (expr->type == NodeType::vstore) {
    emit_code("{}.store(access_{}(context.buffers[0] {}));", expr[1]->var_name,
              expr[0]->new_addresses(0)->node_type_name,
              address_elements(expr[0]->new_addresses(0), "0", 2));
  } else if (expr->type == NodeType::load) {
    emit_code("auto {} = {}::load({});", expr->var_name,
              vv_type(expr->data_type), expr[0]->var_name);
#if (0)
    bool regular = false;
    if (expr[0]->type == NodeType::pointer &&
        expr[0][1]->type == NodeType::index) {
      regular = true;
    } else {
    }
    if (regular) {
      /*
      emit_code("auto {} = load<{}, {}>({}_base, {}_offsets);", expr->var_name,
                vv_width, data_type_name(expr->data_type), expr[0]->var_name,
                expr[0]->var_name);
                */
      // TC_P(expr->members.size());

      std::vector<int> offsets;
      for (int i = 0; i + 1 < (int)expr->members.size(); i++) {
        TC_ASSERT(
            expr->members[i]->addr().same_type(expr->members[i + 1]->addr()));
      }
      for (int i = 0; i < (int)expr->members.size(); i++) {
        offsets.push_back(expr->members[i]->addr().offset());
{'fraction_non_alphanumeric': 0.22717749176893146, 'fraction_numerical': 0.011573381223186671, 'mean_word_length': 2.80128934395146, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 219, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7367821', 'n_tokens_mistral': 3687, 'n_tokens_neox': 3609, 'n_words': 588}