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