diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.cc b/tensorflow/compiler/xla/service/hlo_creation_utils.cc index 846b9cfbeb5..dd174772c62 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.cc +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc @@ -33,6 +33,15 @@ limitations under the License. namespace xla { using absl::StrCat; +StatusOr MakeUnaryHlo(HloOpcode opcode, + HloInstruction* operand) { + HloComputation* computation = operand->parent(); + TF_ASSIGN_OR_RETURN(Shape unary_op_shape, + ShapeInference::InferUnaryOpShape(opcode, operand)); + return computation->AddInstruction( + HloInstruction::CreateUnary(unary_op_shape, opcode, operand)); +} + StatusOr MakeBinaryHlo(HloOpcode opcode, HloInstruction* lhs, HloInstruction* rhs) { HloComputation* computation = lhs->parent(); @@ -344,6 +353,15 @@ StatusOr MakeReduceHlo(HloInstruction* operand, scalar_shape, operand, init_value, all_dims, reduce_computation)); } +StatusOr MakeReverseHlo(HloInstruction* operand, + absl::Span dimensions) { + HloComputation* computation = operand->parent(); + TF_ASSIGN_OR_RETURN(Shape reverse_shape, ShapeInference::InferReverseShape( + operand->shape(), dimensions)); + return computation->AddInstruction( + HloInstruction::CreateReverse(reverse_shape, operand, dimensions)); +} + StatusOr MakeSelectHlo(HloInstruction* pred, HloInstruction* on_true, HloInstruction* on_false, diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.h b/tensorflow/compiler/xla/service/hlo_creation_utils.h index 754f7e2be33..3f2e3aa25a1 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.h +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.h @@ -27,6 +27,11 @@ namespace xla { // ergonomic. We don't have a complete set of helpers yet -- I expect we'll // expand this interface as needed on an ad-hoc basis. +// Creates a unary HLO instruction and adds it to the computation containing +// `operand`. +StatusOr MakeUnaryHlo(HloOpcode opcode, + HloInstruction* operand); + // Creates a binary HLO instruction and adds it to the computation containing // `lhs` and `rhs` (`lhs` and `rhs` must be in the same computation). StatusOr MakeBinaryHlo(HloOpcode opcode, HloInstruction* lhs, @@ -145,6 +150,11 @@ StatusOr MakeReduceHlo(HloInstruction* operand, HloOpcode binary_opcode, HloModule* module); +// Creates a Reverse HLO instruction and adds it to the computation containing +// `operand`. +StatusOr MakeReverseHlo(HloInstruction* operand, + absl::Span dimensions); + // Creates a Select HLO instruction and adds it to the computation containing // the predicate. The on_true and on_false instructions must also be contained // in the same computation. If on_true and on_false are tuples, create a tuple