torch-mlir/projects/ltc/csrc/base_lazy_backend/mlir_node.cpp

//===- mlir_node.cpp ------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// Also available under a BSD-style license. See LICENSE.
//
//===----------------------------------------------------------------------===//
// This file is adapted from pytorch/pytorch
// https://github.com/pytorch/pytorch/blob/lazy_tensor_staging/torch/csrc/lazy/ts_backend/ts_node.cpp
//===----------------------------------------------------------------------===//

#include "mlir_node.h"
#include "utils/exception.h"

namespace torch {
namespace lazy {

namespace {

hash_t OperandHashes(
    const OpList& operands, const c10::ArrayRef<Shape>& shapes,
    const hash_t& seed, bool bakeInSizes) {
  hash_t hash = seed;
  for (auto& operand : operands) {
    if (!operand) {
      hash = HashCombine(hash, static_cast<uint64_t>(kNullOpt));
      continue;
    }
    auto operand_hash = bakeInSizes ? operand.shapeHash() : operand.hash();
    hash = HashCombine(hash, operand_hash);
  }
  for (auto& shape : shapes) {
    hash = HashCombine(hash, shape.hash(bakeInSizes));
  }
  return hash;
}

} // namespace


// Adds a static hook that is run after every single TorchMlirNode is initialized
static std::vector<std::function<void(TorchMlirNode*)>> constructor_hooks;
void TorchMlirNode::addConstructorHook(std::function<void(TorchMlirNode*)> f) {
  constructor_hooks.emplace_back(f);
}

TorchMlirNode::TorchMlirNode(
    OpKind op, OpList operands, std::vector<Shape>&& shapes, size_t num_outputs,
    hash_t hash_seed)
    : Node(op, operands, std::move(shapes), num_outputs) {
  hash_seed = HashCombine(op.hash(), hash_seed);
  shape_hash_ = OperandHashes(operands, this->shapes(), hash_seed, true);
  dag_hash_ =
      (enableDynamicShape()
           ? OperandHashes(operands, this->shapes(), hash_seed, false)
           : shape_hash_);

  for (std::function<void(TorchMlirNode*)>& f : constructor_hooks) {
    f(this);
  }
}

TorchMlirNode::TorchMlirNode(
    OpKind op, OpList operands, const std::function<Shape()>& shape_fn,
    size_t num_outputs, hash_t hash_seed)
    : TorchMlirNode(
          op, operands, std::vector<Shape>{}, num_outputs, hash_seed) {
  addComputedShape(shape_fn);
}

TorchMlirNode::TorchMlirNode(
    OpKind op, OpList operands, size_t num_outputs, hash_t hash_seed)
    : TorchMlirNode(
          op, operands, std::vector<Shape>{}, num_outputs, hash_seed) {}

TorchMlirNode::TorchMlirNode(
    OpKind op, Shape shape, size_t num_outputs, hash_t hash_seed)
    : TorchMlirNode(op, {}, {std::move(shape)}, num_outputs, hash_seed) {}

hash_t TorchMlirNode::hash() const { return dag_hash_; }

hash_t TorchMlirNode::shapeHash() const { return shape_hash_; }


TorchMlirNode* TorchMlirNode::mlir_node(int index) const {
  return dynamic_cast<TorchMlirNode*>(operands_.at(index).get());
}

///////////////////////////////////////////////////////////////////////////////
// TorchMlirTensorList
///////////////////////////////////////////////////////////////////////////////

OpKind TorchMlirTensorList::ClassOpKind() {
  // Note: this OpKind is separate from ltc_ops.h since it would be a circular
  // import otherwise
  static const OpKind tensor_list_opkind =
      OpKind::Get("lazy_tensors::tensor_list");
  return tensor_list_opkind;
}

TorchMlirTensorList::TorchMlirTensorList(OpList values)
    : TorchMlirNode(
          /*op=*/TorchMlirTensorList::ClassOpKind(),
          /*operands=*/values,
          /*shapes=*/std::vector<Shape>(),
          /*num_outputs=*/1,
          /*hash_seed=*/kHashSeed) {}

torch::lazy::TorchMlirOpVector TorchMlirTensorList::Lower(
    TorchMlirFunction function, TorchMlirLoweringContext* loctx) const {
  std::vector<torch::jit::Value*> tensor_list;
  CHECK(!operands().empty());
  for (const torch::lazy::Output& operand : operands()) {
    tensor_list.emplace_back(loctx->GetOutputOp(operand));
  }
  auto graph = function->graph();
  auto listnode =
      graph->insertNode(graph->createList(c10::TensorType::get(), tensor_list));
  return {listnode->output()};
}

///////////////////////////////////////////////////////////////////////////////
// TorchMlirOptionalTensorList
///////////////////////////////////////////////////////////////////////////////

OpKind TorchMlirOptionalTensorList::ClassOpKind() {
  // Note: this OpKind is separate from ltc_ops.h since it would be a circular
  // import otherwise
  static const OpKind tensor_list_opkind =
      OpKind::Get("lazy_tensors::optional_tensor_list");
  return tensor_list_opkind;
}

TorchMlirOptionalTensorList::TorchMlirOptionalTensorList(OpList values)
    : TorchMlirNode(
          /*op=*/TorchMlirOptionalTensorList::ClassOpKind(),
          /*operands=*/values,
          /*shapes=*/std::vector<Shape>(),
          /*num_outputs=*/1,
          /*hash_seed=*/kHashSeed) {}

torch::lazy::TorchMlirOpVector TorchMlirOptionalTensorList::Lower(
    TorchMlirFunction function, TorchMlirLoweringContext* loctx) const {
  std::vector<torch::jit::Value*> tensor_list;
  CHECK(!operands().empty());
  for (const torch::lazy::Output& operand : operands()) {
    tensor_list.emplace_back(loctx->GetOutputOp(operand));
  }
  auto graph = function->graph();
  auto listnode =
      graph->insertNode(graph->createList(c10::OptionalType::create(c10::TensorType::get()), tensor_list));
  return {listnode->output()};
}

} // namespace lazy
} // namespace torch