SymIntOps.cpp

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <executorch/backends/vulkan/runtime/graph/ops/ExecuteNode.h>
#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>

namespace vkcompute {

//
// sym_size
//

void sym_size_impl(ComputeGraph* graph, const std::vector<ValueRef>& args) {
  const ValueRef in_tensor = args.at(0);
  const ValueRef dim = args.at(1);
  const ValueRef out_symint = args.at(2);

  const int64_t dim_val = graph->extract_scalar<int64_t>(dim);
  const int64_t size_at_dim = graph->size_at<int64_t>(dim_val, in_tensor);

  graph->set_symint(out_symint, static_cast<int32_t>(size_at_dim));
}

void resize_sym_size_node(
    ComputeGraph* graph,
    const std::vector<ArgGroup>& args,
    const std::vector<ValueRef>& resize_args) {
  (void)args; // Unused parameter
  sym_size_impl(graph, resize_args);
}

/*
 * This operator takes a tensor and an integer dimension as inputs, and produces
 * a symint as output. The symint's value is the size of the tensor at the
 * specified dimension.
 */
void sym_size_int(ComputeGraph& graph, const std::vector<ValueRef>& args) {
  sym_size_impl(&graph, args);

  graph.execute_nodes().emplace_back(
      new ExecuteNode(resize_sym_size_node, args));
}

//
// binary operators
//

void sym_add_impl(ComputeGraph* graph, const std::vector<ValueRef>& args) {
  const ValueRef a = args.at(0);
  const ValueRef b = args.at(1);
  const ValueRef out = args.at(2);

  const int32_t a_val = graph->read_symint(a);
  const int32_t b_val = graph->read_symint(b);
  const int32_t result = a_val + b_val;

  graph->set_symint(out, result);
}

void resize_sym_add_node(
    ComputeGraph* graph,
    const std::vector<ArgGroup>& args,
    const std::vector<ValueRef>& resize_args) {
  (void)args; // Unused parameter
  sym_add_impl(graph, resize_args);
}

/*
 * This operator takes two symints as inputs and produces a symint as output.
 * The output symint's value is the sum of the two input symints.
 */
void sym_add(ComputeGraph& graph, const std::vector<ValueRef>& args) {
  sym_add_impl(&graph, args);

  graph.execute_nodes().emplace_back(
      new ExecuteNode(resize_sym_add_node, args));
}

void sym_sub_impl(ComputeGraph* graph, const std::vector<ValueRef>& args) {
  const ValueRef a = args.at(0);
  const ValueRef b = args.at(1);
  const ValueRef out = args.at(2);

  const int32_t a_val = graph->read_symint(a);
  const int32_t b_val = graph->read_symint(b);
  const int32_t result = a_val - b_val;

  graph->set_symint(out, result);
}

void resize_sym_sub_node(
    ComputeGraph* graph,
    const std::vector<ArgGroup>& args,
    const std::vector<ValueRef>& resize_args) {
  (void)args;
  sym_sub_impl(graph, resize_args);
}

void sym_sub(ComputeGraph& graph, const std::vector<ValueRef>& args) {
  sym_sub_impl(&graph, args);

  graph.execute_nodes().emplace_back(
      new ExecuteNode(resize_sym_sub_node, args));
}

void sym_floordiv_impl(ComputeGraph* graph, const std::vector<ValueRef>& args) {
  const ValueRef a = args.at(0);
  const ValueRef b = args.at(1);
  const ValueRef out = args.at(2);

  const int32_t a_val = graph->read_symint(a);
  const int32_t b_val = graph->read_symint(b);
  // Floor division: round towards negative infinity
  const int32_t result = (a_val ^ b_val) < 0 && a_val % b_val != 0
      ? a_val / b_val - 1
      : a_val / b_val;

  graph->set_symint(out, result);
}

void resize_sym_floordiv_node(
    ComputeGraph* graph,
    const std::vector<ArgGroup>& args,
    const std::vector<ValueRef>& resize_args) {
  (void)args;
  sym_floordiv_impl(graph, resize_args);
}

void sym_floordiv(ComputeGraph& graph, const std::vector<ValueRef>& args) {
  sym_floordiv_impl(&graph, args);

  graph.execute_nodes().emplace_back(
      new ExecuteNode(resize_sym_floordiv_node, args));
}

void sym_mul_impl(ComputeGraph* graph, const std::vector<ValueRef>& args) {
  const ValueRef a = args.at(0);
  const ValueRef b = args.at(1);
  const ValueRef out = args.at(2);

  const int32_t a_val = graph->read_symint(a);
  const int32_t b_val = graph->read_symint(b);
  const int32_t result = a_val * b_val;

  graph->set_symint(out, result);
}

void resize_sym_mul_node(
    ComputeGraph* graph,
    const std::vector<ArgGroup>& args,
    const std::vector<ValueRef>& resize_args) {
  (void)args;
  sym_mul_impl(graph, resize_args);
}

void sym_mul(ComputeGraph& graph, const std::vector<ValueRef>& args) {
  sym_mul_impl(&graph, args);

  graph.execute_nodes().emplace_back(
      new ExecuteNode(resize_sym_mul_node, args));
}

void select_as_symint_impl(
    ComputeGraph* graph,
    const std::vector<ArgGroup>& unused,
    const std::vector<ValueRef>& args) {
  (void)unused; // Unused parameter

  const ValueRef x = args.at(0);
  const ValueRef dim = args.at(1);
  const ValueRef index = args.at(2);
  const ValueRef out = args.at(3);

  const int64_t dim_val = graph->extract_scalar<int64_t>(dim);
  int64_t index_val = graph->extract_scalar<int64_t>(index);

  const std::vector<int64_t> x_sizes = graph->sizes_of(x);
  const vkapi::ScalarType x_dtype = graph->dtype_of(x);

  if (index_val < 0) {
    index_val += x_sizes[dim_val];
  }

  const StagingPtr x_staging = graph->get_staging(graph->staging_of(x));

  int32_t x_val;
  switch (x_dtype) {
    case vkapi::ScalarType::Int:
      x_val = x_staging->select_element_at_dim<int32_t>(
          x_sizes, dim_val, index_val);
      break;
    case vkapi::ScalarType::Long:
      x_val = static_cast<int32_t>(x_staging->select_element_at_dim<int64_t>(
          x_sizes, dim_val, index_val));
      break;
    default:
      VK_THROW("Unsupported dtype for select_as_symint");
  }

  graph->set_symint(out, x_val);
}

void select_as_symint(ComputeGraph& graph, const std::vector<ValueRef>& args) {
  select_as_symint_impl(&graph, {}, args);

  graph.execute_nodes().emplace_back(new ExecuteNode(
      select_as_symint_impl, args, {}, "select_as_symint", true));
  graph.set_has_data_dependent_shapes();
}

REGISTER_OPERATORS {
  VK_REGISTER_OP(sym_size.int, sym_size_int);
  VK_REGISTER_OP(add, sym_add);
  VK_REGISTER_OP(sub, sym_sub);
  VK_REGISTER_OP(floordiv, sym_floordiv);
  VK_REGISTER_OP(mul, sym_mul);
  VK_REGISTER_OP(et_vk.select_as_symint.default, select_as_symint);
}

} // namespace vkcompute