/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#include "glue_ops.h"
#include <ATen/FunctionalTensorWrapper.h>

void copy_blocks__functionalization_glue(
    const std::vector<torch::Tensor> &k_caches,
    const std::vector<torch::Tensor> &v_caches,
    const c10::Dict<int64_t, c10::List<int64_t>> &block_mapping) {
  TORCH_INTERNAL_ASSERT(at::functionalization::impl::isFunctionalTensor(k_caches));
  TORCH_INTERNAL_ASSERT(at::functionalization::impl::isFunctionalTensor(v_caches));
  at::functionalization::impl::sync(k_caches);
  at::functionalization::impl::sync(v_caches);
  auto new_k_caches = at::functionalization::impl::from_functional_tensor(k_caches);
  auto new_v_caches = at::functionalization::impl::from_functional_tensor(v_caches);

  // Grab the dispatcher entry corresponding to the out-of-place op, "foo"
  static auto op_handle =
      c10::Dispatcher::singleton()
          // specify namespace::op_name, op_overload_name
          .findSchemaOrThrow("torch_mlu_ops::copy_blocks_out_of_place", "")
          // Specify the C++ schema of the out-of-place op.
          .typed<std::tuple<std::vector<at::Tensor>, std::vector<at::Tensor>>(
              const std::vector<at::Tensor> &k_caches, const std::vector<at::Tensor> &v_caches,
              const c10::Dict<int64_t, c10::List<int64_t>> &block_mapping)>();
  // Next, redispatch to the out-of-place op, foo() (user called foo_, we call fooo)
  std::tuple<std::vector<at::Tensor>, std::vector<at::Tensor>> tmp_output;
  {
    at::AutoDispatchSkipFunctionalize guard;
    tmp_output = op_handle.call(new_k_caches, new_v_caches, block_mapping);
  }
  // Finally, tell functionalization about this mutation.
  at::functionalization::impl::replace_(k_caches, std::get<0>(tmp_output));
  at::functionalization::impl::replace_(v_caches, std::get<1>(tmp_output));
  at::functionalization::impl::commit_update(k_caches);
  at::functionalization::impl::commit_update(v_caches);
}
