Home | History | Annotate | Download | only in llvm_ir
      1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LOOP_EMITTER_H_
     17 #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LOOP_EMITTER_H_
     18 
     19 #include <functional>
     20 
     21 #include "llvm/IR/BasicBlock.h"
     22 #include "llvm/IR/IRBuilder.h"
     23 #include "llvm/IR/Value.h"
     24 #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h"
     25 #include "tensorflow/compiler/xla/statusor.h"
     26 
     27 namespace xla {
     28 namespace llvm_ir {
     29 
     30 // A function type for emitting code that generates an element in the target
     31 // array. The function gets a multi-dimensional index as its only input. This
     32 // index specifies the target element for which a value needs to be computed.
     33 // The function has to emit code to compute this value and return the resulting
     34 // llvm::Value*.
     35 using ElementGenerator =
     36     std::function<StatusOr<llvm::Value*>(const IrArray::Index& index)>;
     37 
     38 // Emits a loop for every element in the given shape.
     39 class LoopEmitter {
     40  public:
     41   using BodyEmitter =
     42       std::function<tensorflow::Status(const IrArray::Index& index)>;
     43 
     44   LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape,
     45               llvm::IRBuilder<>* ir_builder);
     46   // Constructs a LoopEmitter from an element generator that generates each
     47   // element of the given target array.
     48   LoopEmitter(const ElementGenerator& target_element_generator,
     49               const IrArray& target_array, llvm::IRBuilder<>* ir_builder);
     50 
     51   // Constructs a LoopEmitter that emits one element into each of N separate
     52   // arrays on each iteration of the loop.
     53   //
     54   // This is used for multi-output fusion.  target_element_generator must
     55   // produce an LLVM struct with N elements.
     56   LoopEmitter(const ElementGenerator& target_element_generator,
     57               tensorflow::gtl::ArraySlice<IrArray> target_arrays,
     58               llvm::IRBuilder<>* ir_builder);
     59 
     60   LoopEmitter(const LoopEmitter&) = delete;
     61   LoopEmitter& operator=(const LoopEmitter&) = delete;
     62   virtual ~LoopEmitter() = default;
     63 
     64   // Emits a loop nest (with a yet-to-be-filled loop body) that iterates through
     65   // every element in the given shape. Returns the multi-dimensional index that
     66   // specifies the element.
     67   IrArray::Index EmitIndexAndSetExitBasicBlock() {
     68     return EmitIndexAndSetExitBasicBlock(/*loop_name=*/"");
     69   }
     70   virtual IrArray::Index EmitIndexAndSetExitBasicBlock(
     71       tensorflow::StringPiece loop_name);
     72 
     73   // Emits a complete loop nest for every element in the given shape.
     74   tensorflow::Status EmitLoop(tensorflow::StringPiece loop_name = "");
     75 
     76  protected:
     77   // An IR emitter that generates the loop body.
     78   BodyEmitter body_emitter_;
     79 
     80   // The shape that the emitted loop iterates through.
     81   Shape shape_;
     82 
     83   // Points to the exit block of the emitted loop. If the given shape is
     84   // scalar, no loops are emitted and exit_bb_ is nullptr in that case.
     85   llvm::BasicBlock* exit_bb_;
     86 
     87   llvm::IRBuilder<>* ir_builder_;
     88 };
     89 
     90 }  // namespace llvm_ir
     91 }  // namespace xla
     92 
     93 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LOOP_EMITTER_H_
     94