message: Fix ordering out output of sort operations
diff: #include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
  #include "tensorflow/compiler/xla/service/hlo_instructions.h"
  +#include "tensorflow/compiler/xla/service/hlo_query.h"
  #include <popops/ElementWise.hpp>
  #include <popops/Sort.hpp>
  namespace xla {
  namespace poplarplugin {
  +namespace {
  +bool IsSimpleComparison(const HloInstruction* inst) {
  + HloInstruction* root(inst->to_apply()->root_instruction());
  + if (!hlo_query::AllOperandsAreParameters(*root)) {
  + return false;
  + }
  +
  + if (root->opcode() != HloOpcode::kCompare) {
  + return false;
  + }
  +
  + switch (root->comparison_direction()) {
  + case ComparisonDirection::kGe:
  + case ComparisonDirection::kGt:
  + case ComparisonDirection::kLe:
  + case ComparisonDirection::kLt:
  + return true;
  + default:
  + return false;
  + }
  +}
  +
  +bool ReverseSortOutput(const HloInstruction* inst) {
  + HloInstruction* root(inst->to_apply()->root_instruction());
  + switch (root->comparison_direction()) {
  + case ComparisonDirection::kGe:
  + case ComparisonDirection::kGt:
  + return true;
  + default:
  + return false;
  + }
  +}
  +}
  +
  StatusOr<poplar::program::Program> CreateSort(CompilerResources& res,
  const HloInstruction* inst,
  TensorMap& tensor_map) {
  const HloSortInstruction* sort = Cast<HloSortInstruction>(inst);
  + if (!IsSimpleComparison(inst)) {
  + return xla::Unimplemented(
  + "Current Sort implementation only supports GT/LT/GE/LE comparisons");
  + }
  +
  poplar::Graph& graph = GetGraph(res, inst);
  poplar::program::Sequence prog;
  @@ -33,6 +73,10 @@ StatusOr<poplar::program::Program> CreateSort(CompilerResources& res,
  popops::sortInPlace(graph, to_sort, sort->dimensions(0), prog);
  + if (ReverseSortOutput(inst)) {
  + TF_ASSIGN_OR_RETURN(to_sort, ReverseTensor(to_sort, sort->dimensions()));
  + }
  +
  TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, to_sort));
  } else if (sort->operand_count() == 2) {
  CHECK_EQ(inputs.size(), 2);
  @@ -43,6 +87,11 @@ StatusOr<poplar::program::Program> CreateSort(CompilerResources& res,
  popops::sortKeyValueInPlace(graph, key, value, sort->dimensions(0), prog);
  + if (ReverseSortOutput(inst)) {
  + TF_ASSIGN_OR_RETURN(key, ReverseTensor(key, sort->dimensions()));
  + TF_ASSIGN_OR_RETURN(value, ReverseTensor(value, sort->dimensions()));
  + }
  +
  TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, key));
  TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 1, value));
  