reduction.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#ifndef TVM_TOPI_REDUCTION_H_
#define TVM_TOPI_REDUCTION_H_
 
#include <tvm/te/operation.h>
#include <tvm/topi/broadcast.h>
#include <tvm/topi/detail/constant_utils.h>
#include <tvm/topi/detail/ravel_unravel.h>
#include <tvm/topi/elemwise.h>
#include <tvm/topi/tags.h>
#include <tvm/topi/transform.h>
 
#include <algorithm>
#include <iterator>
#include <string>
#include <vector>
 
namespace tvm {
namespace topi {
 
using namespace tvm::te;
 
using FReduce = std::function<PrimExpr(PrimExpr source, const ffi::Array<IterVar>& axis,
                                       ffi::Array<PrimExpr> init, Span span)>;
 
using FCommReduce = std::function<ffi::Array<PrimExpr>(
    ffi::Array<PrimExpr> exprs, const ffi::Array<IterVar>& axis, PrimExpr* condition)>;
 
inline std::vector<int> GetRealAxis(int ndim, const ffi::Optional<ffi::Array<int64_t>>& axis) {
  std::vector<int> real_axis;
  if (!axis.has_value()) {
    for (int i = 0; i < ndim; ++i) {
      real_axis.push_back(i);
    }
  } else {
    // Use a set so duplicates are removed and the dims are sorted
    for (int64_t elem : axis.value()) {
      int64_t val = elem;
      if (val < 0) {
        val += ndim;
      }
      TVM_FFI_ICHECK_LT(val, ndim) << " exceeds the maximum dimension " << ndim;
      TVM_FFI_ICHECK_GE(val, 0);
      real_axis.push_back(static_cast<int>(val));
    }
    std::sort(real_axis.begin(), real_axis.end());
    real_axis.resize(std::unique(real_axis.begin(), real_axis.end()) - real_axis.begin());
  }
  return real_axis;
}
 
inline ffi::Array<IterVar> MakeReduceAxes(const std::vector<int>& real_axis, const Tensor& data) {
  ffi::Array<IterVar> reduce_axes;
  for (auto i : real_axis) {
    std::string name = "k" + std::to_string(i);
    reduce_axes.push_back(tvm::te::reduce_axis(Range(0, data->shape[i]), name));
  }
  return reduce_axes;
}
 
inline ffi::Array<PrimExpr> MakeReduceTargetShape(const std::vector<int>& real_axis,
                                                  const Tensor& data, bool keepdims,
                                                  bool atleast1d) {
  auto ndim = data->shape.size();
  ffi::Array<PrimExpr> target_shape;
  if (keepdims) {
    for (size_t i = 0; i < ndim; ++i) {
      if (std::find(real_axis.begin(), real_axis.end(), i) != real_axis.end()) {
        // real_axis contains i
        target_shape.push_back(1);
      } else {
        target_shape.push_back(data->shape[i]);
      }
    }
  } else {
    for (size_t i = 0; i < ndim; ++i) {
      if (std::find(real_axis.begin(), real_axis.end(), i) == real_axis.end()) {
        // real_axis does not contain i
        target_shape.push_back(data->shape[i]);
      }
    }
  }
  if (target_shape.size() == 0 && atleast1d) {
    target_shape.push_back(1);
  }
  return target_shape;
}
 
inline Tensor DoCommReduce(const Tensor& data, FReduce func,
                           const ffi::Array<PrimExpr>& target_shape,
                           const std::vector<int>& reduce_axes,
                           const std::vector<int>& squeeze_axes, Span span = Span()) {
  auto r_axes = MakeReduceAxes(reduce_axes, data);
  auto compute = [&](const ffi::Array<PrimVar>& indices) {
    ffi::Array<PrimExpr> eval_range;
    ffi::Array<PrimVar> eval_indices;
    int arg_counter = 0;
    int red_counter = 0;
 
    for (size_t i = 0; i < data->shape.size(); ++i) {
      bool squeeze_i = std::find(squeeze_axes.begin(), squeeze_axes.end(), i) != squeeze_axes.end();
      if (std::find(reduce_axes.begin(), reduce_axes.end(), i) != reduce_axes.end()) {
        // real_axis contains i
        eval_range.push_back(r_axes[red_counter]);
        eval_indices.push_back(r_axes[red_counter]->var);
        red_counter++;
        arg_counter += !squeeze_i;
        continue;
      }
      eval_range.push_back(indices[arg_counter]);
      arg_counter++;
    }
 
    return func(data(eval_range), r_axes, {}, span);
  };
 
  return tvm::te::compute(target_shape, compute, data->op->name + "_red", kCommReduce);
}
 
inline Tensor CommReduce(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                         FReduce func, bool keepdims, bool atleast1d) {
  auto ndim = data->shape.size();
  if (ndim == 0) {
    auto identity = topi::identity(data, data->op->name + "_red", kCommReduce);
    return atleast1d ? topi::expand_dims(identity, 0, 1) : identity;
  }
  auto real_axis = GetRealAxis(static_cast<int>(ndim), axis);
  auto target_shape = MakeReduceTargetShape(real_axis, data, keepdims, atleast1d);
  return DoCommReduce(data, func, target_shape, real_axis,
                      keepdims ? std::vector<int>() : real_axis);
}
 
inline Tensor CommReduceIdx(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                            FCommReduce func, bool keepdims, bool atleast1d) {
  auto ndim = data->shape.size();
  TVM_FFI_ICHECK_NE(ndim, 0) << "Cannot reduce a 0 dim Tensor";
  auto real_axis = GetRealAxis(static_cast<int>(ndim), axis);
  auto reduce_axes = MakeReduceAxes(real_axis, data);
  auto target_shape = MakeReduceTargetShape(real_axis, data, keepdims, atleast1d);
 
  auto compute = [ndim, keepdims, &real_axis, &reduce_axes, &func,
                  &data](const ffi::Array<PrimVar>& indices) {
    ffi::Array<PrimExpr> eval_range;
    ffi::Array<PrimExpr> eval_indices;
    int arg_counter = 0;
    int red_counter = 0;
 
    for (size_t i = 0; i < ndim; ++i) {
      if (std::find(real_axis.begin(), real_axis.end(), i) != real_axis.end()) {
        // real_axis contains i
        eval_range.push_back(reduce_axes[red_counter]);
        eval_indices.push_back(reduce_axes[red_counter]->var);
        red_counter++;
      } else {
        if (!keepdims) {
          eval_range.push_back(indices[arg_counter]);
          arg_counter++;
        } else {
          eval_range.push_back(indices[i]);
        }
      }
    }
 
    ffi::Array<PrimExpr> ravel_shape;
    for (auto i : real_axis) {
      ravel_shape.push_back(data->shape[i]);
    }
    auto idx = detail::RavelIndex(eval_indices, ravel_shape);
    return func({idx, data(eval_range)}, reduce_axes, nullptr);
  };
 
  auto temp_idx_val =
      tvm::te::compute(target_shape, compute, data->op->name + "_red_temp", kCommReduceIdx);
  auto temp_idx = temp_idx_val[0];
  auto temp_val = temp_idx_val[1];
  return tvm::te::compute(
      target_shape, [&temp_idx](const ffi::Array<PrimVar>& indices) { return temp_idx(indices); },
      data->op->name + "_red", kCommReduceIdx);
}
 
using FCombine =
    std::function<ffi::Array<PrimExpr>(ffi::Array<PrimVar> lhs, ffi::Array<PrimVar> rhs)>;
 
using FIdentity = std::function<ffi::Array<PrimExpr>(std::vector<PrimType> types)>;
 
inline FCommReduce MakeCommReducer(FCombine fcombine, FIdentity fidentity,
                                   std::string name = "reduce") {
  return [fcombine, fidentity, name](ffi::Array<PrimExpr> exprs, const ffi::Array<IterVar>& axis,
                                     PrimExpr* condition) {
    ffi::Array<PrimVar> lhs, rhs;
    ffi::Array<PrimVar> callback_lhs, callback_rhs;
    std::vector<PrimType> dtypes;
 
    for (size_t i = 0; i < exprs.size(); ++i) {
      PrimType dtype = exprs[i].ty();
      dtypes.push_back(dtype);
      PrimVar lhs_var(name + "_lhs_" + std::to_string(i), dtype);
      PrimVar rhs_var(name + "_rhs_" + std::to_string(i), dtype);
      lhs.push_back(lhs_var);
      rhs.push_back(rhs_var);
      callback_lhs.push_back(lhs_var);
      callback_rhs.push_back(rhs_var);
    }
 
    auto result = fcombine(callback_lhs, callback_rhs);
    auto id_elem = fidentity(dtypes);
    auto cond = condition != nullptr ? *condition : IntImm::Bool(true);
 
    auto combiner = tvm::tirx::CommReducer(lhs, rhs, result, id_elem);
    ffi::Array<PrimExpr> outputs;
    for (size_t i = 0; i < exprs.size(); ++i) {
      outputs.push_back(tvm::tirx::Reduce(combiner, exprs, axis, cond, static_cast<int>(i), {}));
    }
    return outputs;
  };
}
 
inline PrimExpr MinOp(PrimExpr source, ffi::Array<IterVar> axis, ffi::Array<PrimExpr> init = {},
                      Span span = Span()) {
  return tvm::min(source, axis, init, span);
}
 
inline PrimExpr MaxOp(PrimExpr source, ffi::Array<IterVar> axis, ffi::Array<PrimExpr> init = {},
                      Span span = Span()) {
  return tvm::max(source, axis, init, span);  // NOLINT(*)
}
 
inline PrimExpr ProdOp(PrimExpr source, ffi::Array<IterVar> axis, ffi::Array<PrimExpr> init = {},
                       Span span = Span()) {
  return tvm::prod(source, axis, init, span);  // NOLINT(*)
}
 
inline Tensor sum(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                  bool keepdims = false, bool atleast1d = false) {
  // Reduction dispatch only depends on boolean element kind; lane encoding is irrelevant here.
  if (data->dtype.code() == DLDataTypeCode::kDLBool) {
    return CommReduce(data, axis, tvm::any, keepdims, atleast1d);
  } else {
    return CommReduce(data, axis, tvm::sum, keepdims, atleast1d);
  }
}
 
inline Tensor collapse_sum(const Tensor& data, ffi::Array<PrimExpr> target_shape) {
  const auto& ishape = data->shape;
  const auto& oshape = target_shape;
  int isize = data->shape.size();
  int osize = target_shape.size();
 
  TVM_FFI_ICHECK_GE(isize, osize)
      << "Invalid collapse: input dimensionality smaller than output dimensionality.\ninput shape: "
      << data->shape << "\nvs\noutput shape: " << target_shape;
 
  std::vector<int> reduce_axes;
  std::vector<int> squeeze_axes;
  tvm::PrimExpr one(1);
 
  for (int i_ax = isize - 1, o_ax = osize - 1; i_ax >= 0; --i_ax) {
    if (o_ax >= 0 && topi::detail::EqualCheck(ishape[i_ax], oshape[o_ax])) {
      --o_ax;
      continue;
    }
    reduce_axes.push_back(i_ax);
    if (o_ax < 0) {  // squeeze o_ax if was added during expansion
      squeeze_axes.push_back(i_ax);
    } else if (topi::detail::EqualCheck(one, oshape[o_ax])) {
      --o_ax;
    }
  }
 
  if (reduce_axes.size() == 0) return topi::identity(data, "tensor", kCommReduce);
 
  std::reverse(reduce_axes.begin(), reduce_axes.end());
  std::reverse(squeeze_axes.begin(), squeeze_axes.end());
  return DoCommReduce(data, tvm::sum, target_shape, reduce_axes, squeeze_axes);
}
 
inline Tensor all(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                  bool keepdims = false, bool atleast1d = false) {
  return CommReduce(data, axis, tvm::all, keepdims, atleast1d);
}
 
inline Tensor any(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                  bool keepdims = false, bool atleast1d = false) {
  return CommReduce(data, axis, tvm::any, keepdims, atleast1d);
}
 
inline Tensor min(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                  bool keepdims = false, bool atleast1d = false) {
  return CommReduce(data, axis, MinOp, keepdims, atleast1d);
}
 
inline Tensor max(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                  bool keepdims = false, bool atleast1d = false) {
  return CommReduce(data, axis, MaxOp, keepdims, atleast1d);
}
 
inline FCommReduce MakeArgminReducer(bool select_last_index = false) {
  // Create a Commutative Reducer with a comparison operation, and method to get the initial value.
  auto fcombine = [=](ffi::Array<PrimVar> lhs, ffi::Array<PrimVar> rhs) {
    ffi::Array<PrimExpr> result;
 
    // Casting to avoid operator ambiguity
    PrimExpr lhs_idx = static_cast<PrimExpr>(lhs[0]);
    PrimExpr rhs_idx = static_cast<PrimExpr>(rhs[0]);
    PrimExpr lhs_val = static_cast<PrimExpr>(lhs[1]);
    PrimExpr rhs_val = static_cast<PrimExpr>(rhs[1]);
 
    // These variables compare the actual values of the array
    auto is_smaller = lhs_val < rhs_val;
    auto is_same = lhs_val == rhs_val;
 
    // This checks if the indices are correct for the reduction. E.g. for select_last_index
    // it gives precedence for later indices of the same element and precedence for sooner
    // indices if not select_last_index;
    PrimExpr proper_index;
    if (select_last_index) {
      proper_index = lhs_idx > rhs_idx;
    } else {
      proper_index = lhs_idx < rhs_idx;
    }
 
    PrimExpr update_index = is_smaller || (is_same && proper_index);
    result.push_back(tvm::tirx::Select(update_index, lhs[0], rhs[0]));  // idx
    result.push_back(tvm::tirx::Select(is_smaller, lhs[1], rhs[1]));    // val
    return result;
  };
  auto fidentity = [&](std::vector<PrimType> types) {
    ffi::Array<PrimExpr> result;
    result.push_back(tvm::tirx::MakeConst(types[0], -1));  // idx
    result.push_back(tvm::max_value(types[1]));            // val
    return result;
  };
  return MakeCommReducer(fcombine, fidentity, "argmin");
}
 
inline Tensor argmin(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                     bool keepdims = false, bool atleast1d = false,
                     bool select_last_index = false) {
  auto reducer = MakeArgminReducer(select_last_index);
  return CommReduceIdx(data, axis, reducer, keepdims, atleast1d);
}
 
inline FCommReduce MakeArgmaxReducer(bool select_last_index = false) {
  // Create a Commutative Reducer with a comparison operation, and method to get the initial value.
  auto fcombine = [=](ffi::Array<PrimVar> lhs, ffi::Array<PrimVar> rhs) {
    ffi::Array<PrimExpr> result;
 
    // Casting to avoid operator ambiguity
    PrimExpr lhs_idx = static_cast<PrimExpr>(lhs[0]);
    PrimExpr rhs_idx = static_cast<PrimExpr>(rhs[0]);
    PrimExpr lhs_val = static_cast<PrimExpr>(lhs[1]);
    PrimExpr rhs_val = static_cast<PrimExpr>(rhs[1]);
 
    // These variables compare the actual values of the array
    auto is_bigger = lhs_val > rhs_val;
    auto is_same = lhs_val == rhs_val;
 
    // This checks if the indices are correct for the reduction. E.g. for select_last_index
    // it gives precedence for later indices of the same element and precedence for sooner
    // indices if not select_last_index;
    PrimExpr proper_index;
    if (select_last_index) {
      proper_index = lhs_idx > rhs_idx;
    } else {
      proper_index = lhs_idx < rhs_idx;
    }
 
    PrimExpr update_index = is_bigger || (is_same && proper_index);
    result.push_back(tvm::tirx::Select(update_index, lhs[0], rhs[0]));  // idx
    result.push_back(tvm::tirx::Select(is_bigger, lhs[1], rhs[1]));     // val
    return result;
  };
  auto fidentity = [&](std::vector<PrimType> types) {
    ffi::Array<PrimExpr> result;
    result.push_back(tvm::tirx::MakeConst(types[0], -1));  // idx
    result.push_back(tvm::min_value(types[1]));            // val
    return result;
  };
  return MakeCommReducer(fcombine, fidentity, "argmax");
}
 
inline Tensor argmax(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                     bool keepdims = false, bool atleast1d = false,
                     bool select_last_index = false) {
  auto reducer = MakeArgmaxReducer(select_last_index);
  return CommReduceIdx(data, axis, reducer, keepdims, atleast1d);
}
 
inline Tensor prod(const Tensor& data, const ffi::Optional<ffi::Array<int64_t>>& axis,
                   bool keepdims = false, bool atleast1d = false) {
  return CommReduce(data, axis, ProdOp, keepdims, atleast1d);
}
 
inline FCommReduce MakeTupleSumReducer() {
  auto fcombine = [](ffi::Array<PrimVar> lhs, ffi::Array<PrimVar> rhs) {
    ffi::Array<PrimExpr> result;
    TVM_FFI_ICHECK_EQ(lhs.size(), rhs.size());
    result.reserve(lhs.size());
    for (size_t i = 0; i < lhs.size(); ++i) {
      result.push_back(lhs[i] + rhs[i]);
    }
    return result;
  };
  auto fidentity = [](std::vector<PrimType> types) {
    ffi::Array<PrimExpr> result;
    for (size_t i = 0; i < types.size(); ++i) {
      result.push_back(tvm::tirx::MakeConst(types[i], 0));
    }
    return result;
  };
  return MakeCommReducer(fcombine, fidentity, "tuple_sum");
}
 
}  // namespace topi
}  // namespace tvm
#endif  // TVM_TOPI_REDUCTION_H_