Function traverse_binary

src/core/tensor/tensor_math_cpp.h:171–230

template <typename DType>
void traverse_binary(const Tensor &in1, const Tensor &in2, Tensor *out,
                     std::function<DType(DType, DType)> func) {
  DType *outPtr = static_cast<DType *>(out->block()->mutable_data());
  const DType *in1Ptr = static_cast<const DType *>(in1.block()->data());
  const DType *in2Ptr = static_cast<const DType *>(in2.block()->data());
  /*
  vector<int> traversal_info_in1 = generate_traversal_info(in1);
  vector<int> traversal_info_in2 = generate_traversal_info(in2);
  vector<int> shape_multipliers_in1 = generate_shape_multipliers(in1);
  vector<int> shape_multipliers_in2 = generate_shape_multipliers(in2);

  for (size_t i = 0; i < in1.Size(); i++) {
    outPtr[i] = func(in1Ptr[traversal_info_in1[in1.shape().size()]],
                     in2Ptr[traversal_info_in2[in2.shape().size()]]);
    traverse_next(in1, shape_multipliers_in1, traversal_info_in1, i + 1);
    traverse_next(in2, shape_multipliers_in2, traversal_info_in2, i + 1);
  }
  */
  auto prod = Product(in1.shape());
  CHECK(in1.shape() == out->shape());
  CHECK(in2.shape() == out->shape());
  if ((in1.stride() == out->stride()) && (in2.stride() == in1.stride())) {
    for (size_t i = 0; i < prod; i++) outPtr[i] = func(in1Ptr[i], in2Ptr[i]);
  } else {
    /*
    LOG(INFO) << "not equal stride";
    std::ostringstream s1, s2, s3, s4, s5, s6;
    std::copy(in1.stride().begin(), in1.stride().end(),
    std::ostream_iterator<int>(s1, ", "));
    std::copy(in2.stride().begin(), in2.stride().end(),
    std::ostream_iterator<int>(s2, ", "));
    std::copy(out->stride().begin(), out->stride().end(),
    std::ostream_iterator<int>(s3, ", "));

    std::copy(in1.shape().begin(), in1.shape().end(),
    std::ostream_iterator<int>(s4, ", "));
    std::copy(in2.shape().begin(), in2.shape().end(),
    std::ostream_iterator<int>(s5, ", "));
    std::copy(out->shape().begin(), out->shape().end(),
    std::ostream_iterator<int>(s6, ", "));

    LOG(INFO) << s1.str() << ": " << s4.str();
    LOG(INFO) << s2.str() << ": " << s5.str();
    LOG(INFO) << s3.str() << ": " << s6.str();
    LOG(INFO) << Backtrace();
    */

    size_t in1_offset = 0, in2_offset = 0, out_offset = 0;
    vector<int> in1_idx(in1.nDim(), 0), in2_idx(in2.nDim(), 0),
        out_idx(out->nDim(), 0);
    for (size_t i = 0; i < prod; i++) {
      outPtr[out_offset] = func(in1Ptr[in1_offset], in2Ptr[in2_offset]);
      out_offset =
          next_offset(out_offset, out->shape(), out->stride(), &out_idx);
      in1_offset = next_offset(in1_offset, in1.shape(), in1.stride(), &in1_idx);
      in2_offset = next_offset(in2_offset, in2.shape(), in2.stride(), &in2_idx);
      // LOG(INFO) <<  in1_offset << ", " << in2_offset << ", " << out_offset;
    }