path: root/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc

/* Copyright 2017 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/cpu/cpu_layout_assignment.h"

#include <initializer_list>
#include <memory>
#include <utility>
#include <vector>

#include "tensorflow/compiler/xla/layout_util.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/service/algebraic_simplifier.h"
#include "tensorflow/compiler/xla/service/computation_layout.h"
#include "tensorflow/compiler/xla/service/cpu/target_machine_features_fake.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_matchers.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/shape_layout.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/test.h"
#include "tensorflow/compiler/xla/test_helpers.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/tests/test_utils.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/array_slice.h"

namespace op = xla::testing::opcode_matchers;

namespace xla {
namespace {

class CpuLayoutAssignmentTest : public HloTestBase {
 protected:
  void AssignLayouts(HloModule* module,
                     ComputationLayout* entry_computation_layout) {
    cpu::TargetMachineFeaturesWithFakeAlignmentLogic target_machine_features(
        [](int64 shape_size) {
          return cpu::TargetMachineFeatures::kEigenExpectedTensorAlignment;
        });
    cpu::CpuLayoutAssignment layout_assignment(entry_computation_layout,
                                               &target_machine_features);
    EXPECT_IS_OK(layout_assignment.Run(module).status());
  }
};

TEST_F(CpuLayoutAssignmentTest, DotWithConstantRhsTensor) {
  auto builder = HloComputation::Builder(TestName());
  Shape lhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 12}, {0, 1});
  Shape rhs_shape = ShapeUtil::MakeShape(F32, {12, 24});
  Shape result_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 24}, {0, 1});
  auto dot_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(0, lhs_shape, "param0"));
  auto dot_rhs = builder.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateFromShape(rhs_shape)));
  auto result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_shape, dot_lhs, dot_rhs));

  auto module = CreateNewModule();
  HloComputation* computation = module->AddEntryComputation(builder.Build());

  ComputationLayout computation_layout(computation->ComputeProgramShape());
  *computation_layout.mutable_parameter_layout(0) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(lhs_shape));
  *computation_layout.mutable_result_layout() =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(result_shape));
  AssignLayouts(module.get(), &computation_layout);

  EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}),
                                dot_lhs->shape().layout()));
  EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({0, 1}),
                                dot_rhs->shape().layout()));
  EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}),
                                result->shape().layout()));
  for (const auto& instruction : computation->instructions()) {
    EXPECT_NE(instruction->opcode(), HloOpcode::kCopy);
  }
}

TEST_F(CpuLayoutAssignmentTest, MultipleDotsWithSameConstantRhsTensor0) {
  // Two dot products have the same constant as the RHS, and both those dot
  // products can be optimized if the constant has a column-major layout.
  auto builder = HloComputation::Builder(TestName());
  Shape lhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 12}, {0, 1});
  Shape rhs_shape = ShapeUtil::MakeShape(F32, {12, 24});
  Shape result_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 24}, {0, 1});
  auto dot_a_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(0, lhs_shape, "param0"));
  auto dot_b_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(1, lhs_shape, "param1"));
  auto dot_rhs = builder.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateFromShape(rhs_shape)));
  auto dot_a_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_shape, dot_a_lhs, dot_rhs));
  auto dot_b_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_shape, dot_b_lhs, dot_rhs));
  builder.AddInstruction(HloInstruction::CreateBinary(
      result_shape, HloOpcode::kAdd, dot_a_result, dot_b_result));

  auto module = CreateNewModule();
  HloComputation* computation = module->AddEntryComputation(builder.Build());

  ComputationLayout computation_layout(computation->ComputeProgramShape());
  *computation_layout.mutable_parameter_layout(0) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(lhs_shape));
  *computation_layout.mutable_result_layout() =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(result_shape));
  AssignLayouts(module.get(), &computation_layout);

  EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({0, 1}),
                                dot_rhs->shape().layout()));
  for (HloInstruction* instruction :
       {dot_a_lhs, dot_b_lhs, dot_a_result, dot_b_result}) {
    EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}),
                                  instruction->shape().layout()));
  }
  for (const auto& instruction : computation->instructions()) {
    EXPECT_NE(instruction->opcode(), HloOpcode::kCopy);
  }
}

TEST_F(CpuLayoutAssignmentTest, MultipleDotsWithSameConstantRhsTensor1) {
  // Two dot products have the same constant as the RHS, but only one of the two
  // dot products can be optimized if the constant has a column-major layout.
  auto builder = HloComputation::Builder(TestName());
  Shape lhs_a_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 12}, {0, 1});
  Shape lhs_b_shape = ShapeUtil::MakeShapeWithLayout(F32, {2, 12}, {0, 1});
  Shape rhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {12, 24}, {0, 1});
  Shape result_a_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 24}, {0, 1});
  Shape result_b_shape = ShapeUtil::MakeShapeWithLayout(F32, {2, 24}, {0, 1});
  auto dot_a_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(0, lhs_a_shape, "param0"));
  auto dot_b_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(1, lhs_b_shape, "param1"));
  auto dot_rhs = builder.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateFromShape(rhs_shape)));
  auto dot_a_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_a_shape, dot_a_lhs, dot_rhs));
  auto dot_b_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_b_shape, dot_b_lhs, dot_rhs));
  auto tuple_result = builder.AddInstruction(
      HloInstruction::CreateTuple({dot_a_result, dot_b_result}));

  auto module = CreateNewModule();
  HloComputation* computation = module->AddEntryComputation(builder.Build());

  ComputationLayout computation_layout(computation->ComputeProgramShape());
  *computation_layout.mutable_parameter_layout(0) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(lhs_a_shape));
  *computation_layout.mutable_parameter_layout(1) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(lhs_b_shape));
  *computation_layout.mutable_result_layout() =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(tuple_result->shape()));
  AssignLayouts(module.get(), &computation_layout);

  for (HloInstruction* instruction :
       {dot_rhs, dot_a_lhs, dot_b_lhs, dot_a_result, dot_b_result}) {
    EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}),
                                  instruction->shape().layout()));
  }
  for (const auto& instruction : computation->instructions()) {
    EXPECT_NE(instruction->opcode(), HloOpcode::kCopy);
  }
}

TEST_F(CpuLayoutAssignmentTest, DotWithConstantLhsTensor) {
  auto builder = HloComputation::Builder(TestName());
  Shape lhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 12}, {0, 1});
  Shape rhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {12, 24}, {0, 1});
  Shape result_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 24}, {0, 1});
  auto dot_lhs = builder.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateFromShape(lhs_shape)));
  auto dot_rhs = builder.AddInstruction(
      HloInstruction::CreateParameter(0, rhs_shape, "param0"));
  auto dot_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_shape, dot_lhs, dot_rhs));

  auto module = CreateNewModule();
  HloComputation* computation = module->AddEntryComputation(builder.Build());

  ComputationLayout computation_layout(computation->ComputeProgramShape());
  *computation_layout.mutable_parameter_layout(0) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(rhs_shape));
  *computation_layout.mutable_result_layout() =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(result_shape));
  AssignLayouts(module.get(), &computation_layout);

  for (HloInstruction* instruction : {dot_lhs, dot_rhs, dot_result}) {
    EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}),
                                  instruction->shape().layout()));
  }
  for (const auto& instruction : computation->instructions()) {
    EXPECT_NE(instruction->opcode(), HloOpcode::kCopy);
  }
}

TEST_F(CpuLayoutAssignmentTest, DotWithConstantRhsTensorThroughGTE) {
  // This is a case we could theoretically optimize at some point, but today we
  // don't.
  auto builder = HloComputation::Builder(TestName());
  Shape lhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 12}, {0, 1});
  Shape rhs_shape = ShapeUtil::MakeShapeWithLayout(F32, {12, 24}, {0, 1});
  Shape other_shape = ShapeUtil::MakeShapeWithLayout(F32, {100, 24}, {0, 1});

  auto constant_shape = ShapeUtil::MakeTupleShape({other_shape, rhs_shape});
  auto constant = builder.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateFromShape(constant_shape)));

  Shape result_shape = ShapeUtil::MakeShape(F32, {1, 24});

  auto dot_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(0, lhs_shape, "param0"));
  auto dot_rhs = builder.AddInstruction(
      HloInstruction::CreateGetTupleElement(rhs_shape, constant, 1));
  auto dot_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(result_shape, dot_lhs, dot_rhs));

  auto module = CreateNewModule();
  HloComputation* computation = module->AddEntryComputation(builder.Build());

  ComputationLayout computation_layout(computation->ComputeProgramShape());
  *computation_layout.mutable_parameter_layout(0) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(lhs_shape));
  *computation_layout.mutable_result_layout() =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(result_shape));
  AssignLayouts(module.get(), &computation_layout);

  for (HloInstruction* instruction : {dot_lhs, dot_rhs, dot_result}) {
    EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}),
                                  instruction->shape().layout()));
  }
  for (const auto& instruction : computation->instructions()) {
    EXPECT_NE(instruction->opcode(), HloOpcode::kCopy);
  }
}

struct DotOutputFusionLayoutAssignmentResult {
  bool layout_assignment_changed_something;
  const HloInstruction* dot_lhs_fusion_param;
  const HloInstruction* dot_rhs_fusion_param;
  const HloInstruction* addend_fusion_param;
};

static StatusOr<DotOutputFusionLayoutAssignmentResult> RunDotOutputFusion(
    HloModule* module, const string& test_name, int m, int k, int n,
    const int64 dot_operand_idx_in_add) {
  DotOutputFusionLayoutAssignmentResult result;

  CHECK(dot_operand_idx_in_add == 0 || dot_operand_idx_in_add == 1);

  auto builder = HloComputation::Builder(test_name);

  Shape dot_lhs_shape = ShapeUtil::MakeShape(F32, {m, k});
  Shape dot_rhs_shape = ShapeUtil::MakeShape(F32, {k, n});
  Shape dot_shape = ShapeUtil::MakeShape(F32, {m, n});

  HloInstruction* dot_lhs = builder.AddInstruction(
      HloInstruction::CreateParameter(0, dot_lhs_shape, "param0"));
  HloInstruction* addend = builder.AddInstruction(
      HloInstruction::CreateParameter(1, dot_shape, "param1"));
  HloInstruction* dot_rhs = builder.AddInstruction(
      HloInstruction::CreateConstant(Literal::CreateFromShape(dot_rhs_shape)));
  HloInstruction* dot_result = builder.AddInstruction(
      HloInstruction::CreateCanonicalDot(dot_shape, dot_lhs, dot_rhs));
  HloInstruction* add_result;
  if (dot_operand_idx_in_add == 0) {
    add_result = builder.AddInstruction(HloInstruction::CreateBinary(
        dot_shape, HloOpcode::kAdd, dot_result, addend));
  } else {
    add_result = builder.AddInstruction(HloInstruction::CreateBinary(
        dot_shape, HloOpcode::kAdd, addend, dot_result));
  }

  HloComputation* computation = module->AddEntryComputation(builder.Build());

  HloInstruction* fusion_instruction =
      module->entry_computation()->AddInstruction(HloInstruction::CreateFusion(
          dot_shape, HloInstruction::FusionKind::kOutput, add_result));
  TF_RETURN_IF_ERROR(
      computation->ReplaceInstruction(add_result, fusion_instruction));

  HloInstruction* fused_add =
      fusion_instruction->fused_instructions_computation()->root_instruction();
  HloInstruction* fused_dot = fusion_instruction->FuseInstruction(dot_result);

  TF_RETURN_IF_ERROR(
      computation->RemoveInstructionAndUnusedOperands(dot_result));

  ComputationLayout computation_layout(computation->ComputeProgramShape());
  *computation_layout.mutable_parameter_layout(0) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(dot_lhs_shape));
  *computation_layout.mutable_parameter_layout(1) =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(dot_shape));
  *computation_layout.mutable_result_layout() =
      ShapeLayout(LayoutUtil::GetWithDefaultLayout(dot_shape));

  result.dot_lhs_fusion_param =
      fusion_instruction->operand(fused_dot->operand(0)->parameter_number());
  result.dot_rhs_fusion_param =
      fusion_instruction->operand(fused_dot->operand(1)->parameter_number());
  result.addend_fusion_param = fusion_instruction->operand(
      fused_add->operand(1 - dot_operand_idx_in_add)->parameter_number());

  cpu::TargetMachineFeaturesWithFakeAlignmentLogic target_machine_features(
      [](int64 shape_size) {
        return cpu::TargetMachineFeatures::kEigenExpectedTensorAlignment;
      });
  cpu::CpuLayoutAssignment layout_assignment(&computation_layout,
                                             &target_machine_features);
  TF_ASSIGN_OR_RETURN(result.layout_assignment_changed_something,
                      layout_assignment.Run(module));

  return result;
}

static void AssertCorrectLayoutForDotOutputFusion(
    const HloComputation* computation,
    const DotOutputFusionLayoutAssignmentResult& layout_assignment_result,
    bool expect_col_major_dot_rhs) {
  Layout expected_dot_rhs_layout = expect_col_major_dot_rhs
                                       ? LayoutUtil::MakeLayout({0, 1})
                                       : LayoutUtil::MakeLayout({1, 0});
  EXPECT_TRUE(LayoutUtil::Equal(
      expected_dot_rhs_layout,
      layout_assignment_result.dot_rhs_fusion_param->shape().layout()));

  EXPECT_TRUE(LayoutUtil::Equal(
      LayoutUtil::MakeLayout({1, 0}),
      layout_assignment_result.dot_lhs_fusion_param->shape().layout()));

  EXPECT_TRUE(LayoutUtil::Equal(
      LayoutUtil::MakeLayout({1, 0}),
      layout_assignment_result.addend_fusion_param->shape().layout()));
  EXPECT_THAT(computation->instructions(), Each(Not(op::Copy())));
}

TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_1x50x19_dot_idx_0) {
  std::unique_ptr<HloModule> module = CreateNewModule();
  TF_ASSERT_OK_AND_ASSIGN(
      DotOutputFusionLayoutAssignmentResult layout_assignment_result,
      RunDotOutputFusion(module.get(), TestName(), /*m=*/1, /*k=*/50, /*n=*/19,
                         /*dot_operand_idx_in_add=*/0));
  ASSERT_TRUE(layout_assignment_result.layout_assignment_changed_something);
  AssertCorrectLayoutForDotOutputFusion(module->entry_computation(),
                                        layout_assignment_result,
                                        /*expect_col_major_dot_rhs=*/true);
}

TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_1x50x19_dot_idx_1) {
  std::unique_ptr<HloModule> module = CreateNewModule();
  TF_ASSERT_OK_AND_ASSIGN(
      DotOutputFusionLayoutAssignmentResult layout_assignment_result,
      RunDotOutputFusion(module.get(), TestName(), /*m=*/1, /*k=*/50, /*n=*/19,
                         /*dot_operand_idx_in_add=*/1));
  ASSERT_TRUE(layout_assignment_result.layout_assignment_changed_something);
  AssertCorrectLayoutForDotOutputFusion(module->entry_computation(),
                                        layout_assignment_result,
                                        /*expect_col_major_dot_rhs=*/true);
}

TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x1_dot_idx_0) {
  std::unique_ptr<HloModule> module = CreateNewModule();
  TF_ASSERT_OK_AND_ASSIGN(
      DotOutputFusionLayoutAssignmentResult layout_assignment_result,
      RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/1,
                         /*dot_operand_idx_in_add=*/0));
  ASSERT_TRUE(layout_assignment_result.layout_assignment_changed_something);
  AssertCorrectLayoutForDotOutputFusion(module->entry_computation(),
                                        layout_assignment_result,
                                        /*expect_col_major_dot_rhs=*/false);
}

TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x1_dot_idx_1) {
  std::unique_ptr<HloModule> module = CreateNewModule();
  TF_ASSERT_OK_AND_ASSIGN(
      DotOutputFusionLayoutAssignmentResult layout_assignment_result,
      RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/1,
                         /*dot_operand_idx_in_add=*/1));
  ASSERT_TRUE(layout_assignment_result.layout_assignment_changed_something);
  AssertCorrectLayoutForDotOutputFusion(module->entry_computation(),
                                        layout_assignment_result,
                                        /*expect_col_major_dot_rhs=*/false);
}

TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x19_dot_idx_0) {
  std::unique_ptr<HloModule> module = CreateNewModule();
  TF_ASSERT_OK_AND_ASSIGN(
      DotOutputFusionLayoutAssignmentResult layout_assignment_result,
      RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/19,
                         /*dot_operand_idx_in_add=*/0));
  ASSERT_TRUE(layout_assignment_result.layout_assignment_changed_something);
  AssertCorrectLayoutForDotOutputFusion(module->entry_computation(),
                                        layout_assignment_result,
                                        /*expect_col_major_dot_rhs=*/false);
}

TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x19_dot_idx_1) {
  std::unique_ptr<HloModule> module = CreateNewModule();
  TF_ASSERT_OK_AND_ASSIGN(
      DotOutputFusionLayoutAssignmentResult layout_assignment_result,
      RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/19,
                         /*dot_operand_idx_in_add=*/1));
  ASSERT_TRUE(layout_assignment_result.layout_assignment_changed_something);
  AssertCorrectLayoutForDotOutputFusion(module->entry_computation(),
                                        layout_assignment_result,
                                        /*expect_col_major_dot_rhs=*/false);
}
}  // namespace
}  // namespace xla