<commit_before>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/horizontal_fusion.h"

#include <algorithm>

#include "absl/container/flat_hash_set.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/layout_util.h"
#include "tensorflow/compiler/xla/service/hlo_creation_utils.h"
#include "tensorflow/core/util/env_var.h"

namespace xla {
namespace gpu {

namespace {

absl::InlinedVector<HloInstruction*, 2> GetOutputsOfFusion(
    const HloInstruction& instr) {
  CHECK(instr.opcode() == HloOpcode::kFusion);
  auto root = instr.fused_expression_root();
  if (root->opcode() != HloOpcode::kTuple) {
    return {root};
  } else {
    return root->operands();
  }
}

// Return the number of outputs of the fused computation.
size_t GetOutputSizeOfFusion(const HloInstruction& instr) {
  CHECK(instr.opcode() == HloOpcode::kFusion);
  auto root = instr.fused_expression_root();
  if (root->opcode() != HloOpcode::kTuple) {
    return 1;
  } else {
    return ShapeUtil::TupleElementCount(root->shape());
  }
}

class HorizontalFusionImpl {
 public:
  HorizontalFusionImpl(HloComputation* computation)
      : computation_(computation) {}

  ~HorizontalFusionImpl() {}

  StatusOr<bool> Run();

 private:
  Status Fuse(absl::Span<HloInstruction*> fused_fusion_instrs);

  // Horizontally fuses `fused_fusion_instrs`. We support only loop fusions
  // and require their numbers of outputs to be the same, because the code
  // generation (of slice input fusion) requires all the concatenated outputs
  // have the same shapes.
  // Returns the fused computation in `uniq_computation` and the operands that
  // are used by `uniq_computation`.
  Status CreateFusedComputation(
      absl::Span<HloInstruction*> fused_fusion_instrs,
      std::unique_ptr<HloComputation>* uniq_computation,
      std::vector<HloInstruction*>* bound_operands);

  // FusionCandidates collects profitable candidates for a given consumer
  // instruction. GetNextSpanOfFusions() can then be iteratively invoked to
  // acquire the next set of fusion candidates based on some heuristics.
  class FusionCandidates {
   public:
    FusionCandidates(HloInstruction* consumer) : fusion_instrs_(), pos_(0) {
      Initialize(consumer);
    }

    // Get a span of fusions to be fused.
    absl::Span<HloInstruction*> GetNextSpanOfFusions();

   private:
    void Initialize(HloInstruction*);

    std::vector<HloInstruction*> fusion_instrs_;
    // pos_ points to the start position of next span.
    size_t pos_;
  };

  HloComputation* computation_;
};  // HorizontalFusionImpl

bool IsFusionSupported(const HloInstruction& instr) {<commit_msg>Polishing comments and coding styles.<commit_after>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/gpu/horizontal_fusion.h"

#include <algorithm>

#include "absl/container/flat_hash_set.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/layout_util.h"
#include "tensorflow/compiler/xla/service/hlo_creation_utils.h"
#include "tensorflow/core/util/env_var.h"

namespace xla {
namespace gpu {

namespace {

absl::InlinedVector<HloInstruction*, 2> GetOutputsOfFusion(
    const HloInstruction& instr) {
  CHECK(instr.opcode() == HloOpcode::kFusion);
  HloInstruction* root = instr.fused_expression_root();
  if (root->opcode() != HloOpcode::kTuple) {
    return {root};
  } else {
    return root->operands();
  }
}

// Returns the number of outputs of the fused computation.
size_t GetOutputSizeOfFusion(const HloInstruction& instr) {
  CHECK(instr.opcode() == HloOpcode::kFusion);
  const HloInstruction* root = instr.fused_expression_root();
  if (root->opcode() != HloOpcode::kTuple) {
    return 1;
  } else {
    return ShapeUtil::TupleElementCount(root->shape());
  }
}

class HorizontalFusionImpl {
 public:
  HorizontalFusionImpl(HloComputation* computation)
      : computation_(computation) {}

  ~HorizontalFusionImpl() {}

  StatusOr<bool> Run();

 private:
  Status Fuse(absl::Span<HloInstruction*> fused_fusion_instrs);

  // Horizontally fuses `fused_fusion_instrs`. It is required that each of
  // `fused_fusion_instrs` is a kLoop fusion. Also, we require their numbers of
  // outputs to be the same, so that each output will be fused/concatenated with
  // the same number of outputs from other fused fusion instrs. Then, all the
  // fused outputs still have the same shapes for kernel generation.
  //
  // Returns the fused computation in `uniq_computation` and the operands that
  // are used by `uniq_computation`.
  Status CreateFusedComputation(
      absl::Span<HloInstruction*> fused_fusion_instrs,
      std::unique_ptr<HloComputation>* uniq_computation,
      std::vector<HloInstruction*>* bound_operands);

  // FusionCandidates collects profitable candidates for a given consumer
  // instruction. GetNextSpanOfFusions() can then be iteratively invoked to
  // acquire the next set of fusion candidates based on some heuristics.
  class FusionCandidates {
   public:
    FusionCandidates(HloInstruction* consumer) : fusion_instrs_(), pos_(0) {
      Initialize(consumer);
    }

    // Gets a span of fusions to be fused.
    absl::Span<HloInstruction*> GetNextSpanOfFusions();

   private:
    void Initialize(HloInstruction*);

    std::vector<HloInstruction*> fusion_instrs_;
    // `pos_` points to the start position of the next span.
    size_t pos_;
  };

  HloComputation* computation_;
};  // HorizontalFusionImpl

bool IsFusionSupported(const HloInstruction& instr) {
{'fraction_non_alphanumeric': 0.11799410029498525, 'fraction_numerical': 0.0032448377581120944, 'mean_word_length': 4.576480263157895, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3629538', 'n_tokens_mistral': 2034, 'n_tokens_neox': 1978, 'n_words': 605}