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_CPU_DOT_OP_EMITTER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_DOT_OP_EMITTER_H_ 18 19 #include "absl/strings/string_view.h" 20 #include "llvm/IR/IRBuilder.h" 21 #include "tensorflow/compiler/xla/service/cpu/cpu_options.h" 22 #include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" 23 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 24 #include "tensorflow/compiler/xla/service/hlo_module_config.h" 25 #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" 26 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" 27 #include "tensorflow/compiler/xla/types.h" 28 #include "tensorflow/core/lib/core/status.h" 29 #include "tensorflow/core/platform/types.h" 30 31 namespace xla { 32 namespace cpu { 33 // Returns true if the two operands and the output of `dot_instr` must have row 34 // major layout. 35 bool DotOperandsAndResultMustHaveRowMajorLayout( 36 const HloInstruction& dot_instr, 37 const TargetMachineFeatures& target_machine_features); 38 39 // Returns true our lowering strategy for `dot_instr` can fold in transposes to 40 // the either of the inputs. 41 bool DotImplementationCanHandleTranspose( 42 const HloInstruction& dot_instr, 43 const TargetMachineFeatures& target_machine_features); 44 45 // Returns the index for an operand to `hlo` that should ideally be column 46 // major. Returns nullopt if there is no such operand or if `hlo` is not a dot 47 // or a fusion containing a dot. 48 absl::optional<int64> ProfitableToMakeDotOperandColumnMajor( 49 const HloInstruction& hlo); 50 51 // Emit LLVM IR to perform the dot operation on lhs_array and rhs_array and 52 // place the result in target_array. IR is emitted at current insert point of 53 // the builder. Upon completion of the method, the insert point is set to the 54 // end of all instructions emitted for this operation. 55 // 56 // If `addend_array` is not nullptr then it must be an array of the same 57 // dimensions as the result, and the result is computed as `addend_array` + 58 // dot(`lhs_array`, `rhs_array`). A non-null `addend_array` is only supported 59 // for Matrix-vector products. 60 Status EmitDotOperation(const HloInstruction& dot, 61 const llvm_ir::IrArray& target_array, 62 const llvm_ir::IrArray& lhs_array, 63 const llvm_ir::IrArray& rhs_array, 64 const llvm_ir::IrArray* addend_array, 65 llvm::Value* executable_run_options_value, 66 llvm::IRBuilder<>* b, 67 const HloModuleConfig& hlo_module_config, 68 const TargetMachineFeatures& target_machine_features); 69 } // namespace cpu 70 } // namespace xla 71 72 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_DOT_OP_EMITTER_H_ 73