Support evaluation in the absence of layouts when possible

Only bitcast requires the layout to be known when evaluating HLO.
In all other cases, we can evaluate without knowing the layout.
This is needed for collective pipelining where we have to analyse while loops before layouts were assigned.

PiperOrigin-RevId: 715163612

Author: frgossen

xla/hlo/evaluator/BUILD

@@ -138,6 +138,7 @@ xla_cc_test(
         "//xla/hlo/analysis:tuple_points_to_analysis",
         "//xla/hlo/builder:xla_builder",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/parser:hlo_parser",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/transforms/simplifiers:hlo_element_type_converter",
         "//xla/service:call_graph",

xla/hlo/evaluator/hlo_evaluator.cc

@@ -40,6 +40,7 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/functional/function_ref.h"
+#include "absl/log/check.h"
 #include "absl/memory/memory.h"
 #include "absl/numeric/bits.h"
 #include "absl/status/status.h"
@@ -899,8 +900,10 @@ absl::StatusOr<Literal> HloEvaluator::Evaluate(
     const auto& computation_shape =
         computation.parameter_instruction(i)->shape();
     const auto& arg_shape = arg_literals[i]->shape();
-    if (!Shape::Equal().MinorToMajorOnlyInLayout()(computation_shape,
-                                                   arg_shape)) {
+    bool ignore_layout = !computation_shape.has_layout();
+    if (!Shape::Equal()
+             .IgnoreLayout(ignore_layout)
+             .MinorToMajorOnlyInLayout()(computation_shape, arg_shape)) {
       return InvalidArgument(
           "Shape mismatch at parameter %d. Computation expected %s, but arg "
           "was %s.",
@@ -1290,12 +1293,30 @@ absl::Status HloEvaluator::EvaluateInternal(
 }
 
 absl::Status HloEvaluator::HandleBitcast(const HloInstruction* bitcast) {
-  const Literal& operand_literal = GetEvaluatedLiteralFor(bitcast->operand(0));
-  Literal result(bitcast->shape());
+  Shape result_shape = bitcast->shape();
+
+  // Allow effective scalars without layouts as the result is unambiguous.
+  if (!result_shape.has_layout() &&
+      ShapeUtil::IsEffectiveScalar(result_shape)) {
+    result_shape = LayoutUtil::GetWithDefaultLayout(result_shape);
+  }
+
+  // In general, we require a layout to evaluate a bitcast: this is the only
+  // operation where indexing is physical rather than logical.
+  if (!result_shape.has_layout()) {
+    return InvalidArgument(
+        "Evaluator cannot evaluate bitcast for non-scalar operand without "
+        "assigned layout.");
+  }
+  TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(result_shape));
+
+  Literal result(result_shape);
+
   // Bitcast output is allowed to be smaller than the input if the backend-
   // specific buffer sizes for the input and output are the same. Since the HLO
   // evaluator doesn't have access to the backend-specific shape size function,
   // assume it's OK to bitcast if output <= input.
+  const Literal& operand_literal = GetEvaluatedLiteralFor(bitcast->operand(0));
   TF_RET_CHECK(operand_literal.size_bytes() >= result.size_bytes());
   memcpy(result.untyped_data(), operand_literal.untyped_data(),
          result.size_bytes());
@@ -1372,8 +1393,11 @@ absl::Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
 #ifndef NDEBUG
     const Literal* input_literal = arg_literals_[parameter->parameter_number()];
     VLOG(2) << "Parameter evaluated to: " << input_literal->ToString();
-    DCHECK(Shape::Equal().MinorToMajorOnlyInLayout()(parameter->shape(),
-                                                     input_literal->shape()))
+    bool check_layout = parameter->shape().has_layout();
+    DCHECK(Shape::Equal()
+               .IgnoreLayout(!check_layout)
+               .MinorToMajorOnlyInLayout()(parameter->shape(),
+                                           input_literal->shape()))
         << "parameter shape is: "
         << ShapeUtil::HumanStringWithLayout(parameter->shape())
         << ", but input literal shape is: "
@@ -4723,7 +4747,7 @@ absl::Status HloEvaluator::Preprocess(const HloInstruction* hlo) {
       }
     }
   }
-  return ShapeUtil::ValidateShape(hlo->shape());
+  return ShapeUtil::ValidateShapeWithOptionalLayout(hlo->shape());
 }
 
 absl::Status HloEvaluator::Postprocess(const HloInstruction* hlo) {

xla/hlo/evaluator/hlo_evaluator_test.cc

@@ -46,6 +46,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/transforms/simplifiers/hlo_element_type_converter.h"
 #include "xla/layout_util.h"
@@ -4566,6 +4567,71 @@ ENTRY main {
   }
 }
 
+TEST_P(HloEvaluatorBf16Test, BitcastWithoutLayout) {
+  const absl::string_view hlo_text_base = R"(
+HloModule Bitcast
+
+ENTRY main {
+  param = %s[2,4] parameter(0)
+  ROOT bitcast = %s[4,2,1] bitcast(%s[2,4] param)
+}
+)";
+  std::string hlo_text;
+  Literal arg;
+  if (use_bfloat16_) {
+    hlo_text = absl::StrFormat(hlo_text_base, "bf16", "bf16", "bf16");
+    arg = LiteralUtil::CreateR2<bfloat16>(
+        {{bfloat16(1), bfloat16(2), bfloat16(3), bfloat16(4)},
+         {bfloat16(5), bfloat16(6), bfloat16(7), bfloat16(8)}});
+  } else {
+    hlo_text = absl::StrFormat(hlo_text_base, "f32", "f32", "f32");
+    arg = LiteralUtil::CreateR2<float>({{1., 2., 3., 4.}, {5., 6., 7., 8.}});
+  }
+
+  HloParserOptions parser_config;
+  parser_config.set_fill_missing_layouts(false);
+  TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnUnverifiedModule(
+                                  hlo_text, HloModuleConfig(), parser_config));
+
+  absl::StatusOr<Literal> actual = Evaluate({&arg});
+  EXPECT_FALSE(actual.ok());
+  EXPECT_EQ(actual.status().message(),
+            "Evaluator cannot evaluate bitcast for non-scalar operand without "
+            "assigned layout.");
+}
+
+TEST_P(HloEvaluatorBf16Test, EffectiveScalarBitcastWithoutLayout) {
+  const absl::string_view hlo_text_base = R"(
+HloModule Bitcast
+
+ENTRY main {
+  param = %s[1,1] parameter(0)
+  ROOT bitcast = %s[1,1,1] bitcast(%s[1,1] param)
+}
+)";
+  std::string hlo_text;
+  Literal arg;
+  if (use_bfloat16_) {
+    hlo_text = absl::StrFormat(hlo_text_base, "bf16", "bf16", "bf16");
+    arg = LiteralUtil::CreateR2<bfloat16>({{bfloat16(2)}});
+  } else {
+    hlo_text = absl::StrFormat(hlo_text_base, "f32", "f32", "f32");
+    arg = LiteralUtil::CreateR2<float>({{2.}});
+  }
+
+  HloParserOptions parser_config;
+  parser_config.set_fill_missing_layouts(false);
+  TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnUnverifiedModule(
+                                  hlo_text, HloModuleConfig(), parser_config));
+
+  TF_ASSERT_OK_AND_ASSIGN(Literal actual, Evaluate({&arg}));
+  if (use_bfloat16_) {
+    EXPECT_TRUE(absl::c_equal(arg.data<bfloat16>(), actual.data<bfloat16>()));
+  } else {
+    EXPECT_TRUE(absl::c_equal(arg.data<float>(), actual.data<float>()));
+  }
+}
+
 // Check that s32 under/overflow doesn't trigger a ubsan failure.
 TEST_F(HloEvaluatorTest, Int32Overflow) {
   const absl::string_view hlo_text = R"(

xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc

@@ -275,12 +275,15 @@ TEST_F(HloConstantFoldingTest, ConstantFoldReduceNoLayout) {
                           ParseAndReturnVerifiedModule(kConstantFoldReduce));
   HloInstruction* add = (*m->computations().begin())->root_instruction();
   LayoutUtil::ClearLayout(add->mutable_shape());
+
   HloConstantFolding const_folder;
   TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(m.get()));
-  EXPECT_FALSE(result);
+  EXPECT_TRUE(result);
 
-  EXPECT_THAT(m->entry_computation()->root_instruction(),
-              GmockMatch(m::Reduce()));
+  EXPECT_EQ(6, m->entry_computation()
+                   ->root_instruction()
+                   ->literal()
+                   .GetFirstElement<int32_t>());
 }
 
 const char* const kConstantFoldLargePad = R"(

xla/literal.cc

@@ -137,7 +137,9 @@ const Shape* TryInternShape(const Shape& shape) {
     return &NilShape();
   }
   if (shape.IsArray() && shape.dimensions_size() == 0 && shape.is_static() &&
-      shape.layout().tiles_size() == 0 && shape.layout().memory_space() == 0) {
+      shape.has_layout() && shape.layout().tiles_size() == 0 &&
+      shape.layout().memory_space() == 0 &&
+      shape.layout().element_size_in_bits() == 0) {
     return &ScalarShape(shape.element_type());
   }
   return nullptr;
@@ -252,18 +254,20 @@ Literal::Literal(const Shape& shape)
     : Literal(shape, /*allocate_arrays=*/true) {}
 
 void Literal::SetShape(const Shape& shape) {
-  Shape shape_storage;
-  const Shape* shape_ptr = &shape;
-  if (shape.IsArray() && LayoutUtil::HasCustomElementSizeInBits(shape)) {
-    shape_storage = shape;
-    shape_storage.mutable_layout()->set_element_size_in_bits(0);
-    shape_ptr = &shape_storage;
-  }
-  if (const Shape* intered_shape_ptr = TryInternShape(*shape_ptr)) {
+  if (const Shape* intered_shape_ptr = TryInternShape(shape)) {
     shape_ = intered_shape_ptr;
-  } else {
-    shape_ = std::make_unique<Shape>(*shape_ptr);
+    return;
+  }
+  auto owning_shape_ptr = std::make_unique<Shape>(shape);
+  if (owning_shape_ptr->IsArray() && !owning_shape_ptr->has_layout()) {
+    *owning_shape_ptr->mutable_layout() =
+        LayoutUtil::GetDefaultLayoutForShape(*owning_shape_ptr);
+  }
+  if (owning_shape_ptr->IsArray() &&
+      LayoutUtil::HasCustomElementSizeInBits(*owning_shape_ptr)) {
+    owning_shape_ptr->mutable_layout()->set_element_size_in_bits(0);
   }
+  shape_ = std::move(owning_shape_ptr);
 }
 
 void Literal::SetPiece(const Shape& shape, Piece* piece, bool allocate_arrays,

xla/service/collective_pipeliner_test.cc

@@ -191,7 +191,7 @@ ENTRY entry {
   EXPECT_EQ(get_tuple_index->tuple_index(), 3);
 }
 
-TEST_F(CollectivePipelinerTest, MinimalCase) {
+TEST_F(CollectivePipelinerTest, MinimalCaseWithoutDefaultLayouts) {
   constexpr absl::string_view hlo_string = R"(
     HloModule module
 
@@ -235,8 +235,10 @@ TEST_F(CollectivePipelinerTest, MinimalCase) {
       ROOT dst_data = bf16[3,8,128] get-tuple-element(while), index=1
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto module,
-                          ParseAndReturnUnverifiedModule(hlo_string, config_));
+  HloParserOptions parser_config;
+  parser_config.set_fill_missing_layouts(false);
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(
+                                           hlo_string, config_, parser_config));
   EXPECT_THAT(RunOptimizer(module.get(), /*last_run=*/true),
               IsOkAndHolds(true));
 

xla/shape.h

@@ -257,8 +257,8 @@ class Shape {
 
     bool operator()(const Shape& lhs, const Shape& rhs);
 
-    Equal& IgnoreLayout() {
-      ignore_layout_ = true;
+    Equal& IgnoreLayout(bool ignore_layout = true) {
+      ignore_layout_ = ignore_layout;
       return *this;
     }
     Equal& IgnoreTilesInLayout() {