xref: /external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc

/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed 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.
==============================================================================*/

#include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h"

#include <numeric>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/core/platform/logging.h"

namespace xla {
namespace llvm_ir {

ForLoop::ForLoop(absl::string_view prefix, absl::string_view suffix,
                 llvm::Value* start_index, llvm::Value* end_index,
                 llvm::Value* step, UnrollMode unroll_mode,
                 bool prevent_vectorization)
    : prefix_(prefix),
      suffix_(suffix),
      start_index_(start_index),
      end_index_(end_index),
      step_(step),
      insert_before_bb_(nullptr),
      unroll_mode_(unroll_mode),
      prevent_vectorization_(prevent_vectorization) {}

/* static */ std::unique_ptr<ForLoop> ForLoop::EmitForLoop(
    absl::string_view prefix, llvm::Value* start_index, llvm::Value* end_index,
    llvm::Value* step, llvm::IRBuilder<>* b, UnrollMode unroll_mode,
    bool prevent_vectorization) {
  std::unique_ptr<ForLoop> loop(new ForLoop(prefix, /*suffix=*/"", start_index,
                                            end_index, step, unroll_mode,
                                            prevent_vectorization));
  loop->Emit(b);
  return loop;
}

void ForLoop::Emit(llvm::IRBuilder<>* b) {
  // The preheader block is the block the builder is currently emitting
  // code into.
  preheader_bb_ = b->GetInsertBlock();

  llvm::BasicBlock::iterator insert_point = b->GetInsertPoint();
  if (insert_point == preheader_bb_->end()) {
    // We're emitting the loop at the end of a basic block. Verify there is no
    // terminator (eg, branch) in the basic block.
    CHECK_EQ(nullptr, preheader_bb_->getTerminator());

    exit_bb_ = CreateLoopBB("loop_exit", b);
  } else {
    // We're emitting the loop into the middle of a basic block. splitBasicBlock
    // requires that this basic block be well-formed (have a terminator).
    CHECK_NE(nullptr, preheader_bb_->getTerminator());

    // Split the preheader to create an exit basic block. The exit basic block
    // will contain all instructions at or after insert_point.
    exit_bb_ = preheader_bb_->splitBasicBlock(insert_point,
                                              GetQualifiedName("loop_exit"));

    // splitBasicBlock adds an unconditional branch between the split basic
    // blocks. Remove it. An unconditional branch will be added below from the
    // preheader to the header.
    preheader_bb_->getTerminator()->eraseFromParent();
  }
  insert_before_bb_ = exit_bb_;

  // Create remaining basic block which form the inside of the loop.
  header_bb_ = CreateLoopBB("loop_header", b);
  body_bb_ = CreateLoopBB("loop_body", b);

  // Function entry basic block.
  // Emit alloca for the induction variable. We do this at the entry to the
  // basic block to ensure the alloc only executes once per function (we could
  // be emitting a nested loop).
  llvm::Function* func = preheader_bb_->getParent();
  b->SetInsertPoint(&func->getEntryBlock(),
                    func->getEntryBlock().getFirstInsertionPt());
  llvm::Value* indvar_address = b->CreateAlloca(
      start_index_->getType(), nullptr, GetQualifiedName("invar_address"));

  // Preheader basic block.
  // Initialize induction variable starting index. Create branch to the header.
  b->SetInsertPoint(preheader_bb_);
  b->CreateStore(start_index_, indvar_address);
  // The preheader should not have a branch yet.
  CHECK_EQ(preheader_bb_->getTerminator(), nullptr);
  b->CreateBr(header_bb_);

  // Header basic block.
  // Emit the loop conditional branch. Load and compare indvar with ending
  // index and jump to loop exit if equal. Jump to body otherwise.
  b->SetInsertPoint(header_bb_);
  indvar_ = b->CreateLoad(start_index_->getType(), indvar_address,
                          GetQualifiedName("indvar"));
  llvm::Value* exit_cond = b->CreateICmpUGE(indvar_, end_index_);
  b->CreateCondBr(/*Cond=*/exit_cond,
                  /*True=*/exit_bb_, /*False=*/body_bb_);

  // Body basic block.
  // Increment indvar, store indvar, and jump to header.
  b->SetInsertPoint(body_bb_);
  llvm::Value* step = step_;
  llvm::Value* indvar = indvar_;

  llvm::Value* indvar_inc = b->CreateAdd(indvar, step, "invar.inc",
                                         /*HasNUW=*/true, /*HasNSW=*/true);
  b->CreateStore(indvar_inc, indvar_address);
  llvm::BranchInst* back_branch = b->CreateBr(header_bb_);

  std::vector<llvm::Metadata*> loop_metadata = GetLoopMetadata(b);
  if (!loop_metadata.empty()) {
    llvm::LLVMContext* ctx = &start_index_->getContext();
    auto temp_node = llvm::MDNode::getTemporary(*ctx, llvm::None);
    loop_metadata.insert(loop_metadata.begin(), temp_node.get());
    auto loop_id = llvm::MDNode::get(*ctx, loop_metadata);
    loop_id->replaceOperandWith(0, loop_id);
    back_branch->setMetadata(llvm::LLVMContext::MD_loop, loop_id);
  }

  // Re-point the IR builder to the loop exit block.
  b->SetInsertPoint(exit_bb_);
}

std::vector<llvm::Metadata*> ForLoop::GetLoopMetadata(llvm::IRBuilder<>* b) {
  const char* const kLlvmLoopUnrollDisableMDName = "llvm.loop.unroll.disable";
  const char* const kLlvmLoopUnrollFullMDName = "llvm.loop.unroll.full";
  const char* const kLlvmLoopVectorizeMDName = "llvm.loop.vectorize.enable";
  llvm::LLVMContext* ctx = &start_index_->getContext();

  std::vector<llvm::Metadata*> result;
  if (unroll_mode_ == xla::llvm_ir::UnrollMode::kNoUnroll) {
    result.push_back(llvm::MDNode::get(
        *ctx, {llvm::MDString::get(*ctx, kLlvmLoopUnrollDisableMDName)}));
  }

  if (prevent_vectorization_) {
    result.push_back(llvm::MDNode::get(
        *ctx, {llvm::MDString::get(*ctx, kLlvmLoopVectorizeMDName),
               llvm::ConstantAsMetadata::get(b->getFalse())}));
  }

  if (unroll_mode_ == xla::llvm_ir::UnrollMode::kFullyUnroll) {
    result.push_back(llvm::MDNode::get(
        *ctx, {llvm::MDString::get(*ctx, kLlvmLoopUnrollFullMDName)}));
  }
  return result;
}

std::string ForLoop::GetQualifiedName(absl::string_view name) {
  return llvm_ir::IrName(prefix_, llvm_ir::IrName(name, suffix_));
}

llvm::BasicBlock* ForLoop::CreateLoopBB(absl::string_view name,
                                        llvm::IRBuilder<>* b) {
  return CreateBasicBlock(insert_before_bb_, GetQualifiedName(name), b);
}

std::unique_ptr<ForLoop> ForLoopNest::AddLoop(absl::string_view suffix,
                                              llvm::Value* start_index,
                                              llvm::Value* end_index,
                                              UnrollMode unroll_mode,
                                              bool prevent_vectorization) {
  return AddLoop(suffix, start_index, end_index, GetConstantWithIndexType(1),
                 unroll_mode, prevent_vectorization);
}

std::unique_ptr<ForLoop> ForLoopNest::AddLoop(
    absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index,
    llvm::Value* stride, UnrollMode unroll_mode, bool prevent_vectorization) {
  if (inner_loop_body_bb_ != nullptr) {
    // Create this loop inside the previous one.
    b_->SetInsertPoint(&*inner_loop_body_bb_->getFirstInsertionPt());
  }
  std::unique_ptr<ForLoop> loop(new ForLoop(
      /*prefix=*/name_, suffix, start_index, end_index, stride, unroll_mode,
      prevent_vectorization));
  loop->Emit(b_);

  if (outer_loop_preheader_bb_ == nullptr) {
    outer_loop_preheader_bb_ = loop->GetPreheaderBasicBlock();
  }

  if (outer_loop_exit_bb_ == nullptr) {
    outer_loop_exit_bb_ = loop->GetExitBasicBlock();
  }

  inner_loop_body_bb_ = loop->GetBodyBasicBlock();

  return loop;
}

std::unique_ptr<ForLoop> ForLoopNest::AddLoop(int64_t start_index,
                                              int64_t end_index,
                                              absl::string_view suffix,
                                              UnrollMode unroll_mode,
                                              bool prevent_vectorization) {
  CHECK_LE(start_index, end_index);
  return AddLoop(suffix, GetConstantWithIndexType(start_index),
                 GetConstantWithIndexType(end_index), unroll_mode,
                 prevent_vectorization);
}

std::unique_ptr<ForLoop> ForLoopNest::AddLoop(int64_t start_index,
                                              int64_t end_index, int64_t stride,
                                              absl::string_view suffix,
                                              UnrollMode unroll_mode,
                                              bool prevent_vectorization) {
  CHECK_LE(start_index, end_index);
  return AddLoop(suffix, GetConstantWithIndexType(start_index),
                 GetConstantWithIndexType(end_index),
                 GetConstantWithIndexType(stride), unroll_mode,
                 prevent_vectorization);
}

IrArray::Index ForLoopNest::AddLoopsForShape(const Shape& shape,
                                             absl::string_view suffix) {
  std::vector<int64_t> dimensions(shape.rank());
  std::iota(dimensions.begin(), dimensions.end(), 0);
  return IrArray::Index(AddLoopsForShapeOnDimensions(shape, dimensions, suffix),
                        shape, index_type_);
}

std::vector<llvm::Value*> ForLoopNest::AddLoopsForShapeOnDimensions(
    const Shape& shape, absl::Span<const int64_t> dimensions,
    absl::string_view suffix) {
  std::vector<llvm::Value*> multi_index(shape.dimensions_size());
  for (int64_t dimension : dimensions) {
    std::unique_ptr<llvm_ir::ForLoop> loop = AddLoop(
        /*start_index=*/0,
        /*end_index=*/shape.dimensions(dimension),
        /*suffix=*/
        llvm_ir::IrName(suffix, absl::StrCat(dimension)));
    multi_index[dimension] = loop->GetIndVarValue();
  }
  return multi_index;
}

std::vector<llvm::Value*> ForLoopNest::EmitOperandArrayLoopNest(
    const llvm_ir::IrArray& operand_array, int64_t dimension_to_skip,
    absl::string_view name_suffix) {
  // Prepares the dimension list we will use to emit the loop nest. Outermost
  // loops are added first. Add loops in major-to-minor order, and skip the
  // 'dimension_to_skip' dimension.
  std::vector<int64_t> dimensions;
  const Shape& shape = operand_array.GetShape();
  // Initially get the dimensions in minor to major order, then reverse them.
  for (int64_t dimension : LayoutUtil::MinorToMajor(shape)) {
    if (dimension != dimension_to_skip) {
      dimensions.push_back(dimension);
    }
  }
  absl::c_reverse(dimensions);

  // Create loop nest with one for-loop for each dimension of the
  // output.
  std::vector<llvm::Value*> multi_index =
      AddLoopsForShapeOnDimensions(shape, dimensions, name_suffix);
  // Verify every dimension except the 'dimension_to_skip' dimension was set in
  // the index.
  for (size_t dimension = 0; dimension < multi_index.size(); ++dimension) {
    if (dimension == dimension_to_skip) {
      DCHECK_EQ(nullptr, multi_index[dimension]);
    } else {
      DCHECK_NE(nullptr, multi_index[dimension]);
    }
  }
  return multi_index;
}

}  // namespace llvm_ir
}  // namespace xla