From c2cc0208151c174cd9e8c469e3167f5cdfd90720 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 4 Apr 2024 12:16:37 -0700
Subject: [PATCH] Add SourceLocation information to xla::Unimplemented errors.

PiperOrigin-RevId: 621934438
---
 xla/service/computation_layout.cc          |  6 ++++--
 xla/service/dynamic_dimension_inference.cc | 19 ++++++++++++-------
 xla/util.h                                 |  2 +-
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/xla/service/computation_layout.cc b/xla/service/computation_layout.cc
index 117b04f3aabf7..675c9635a233e 100644
--- a/xla/service/computation_layout.cc
+++ b/xla/service/computation_layout.cc
@@ -65,7 +65,8 @@ ComputationLayout::FlattenedParameterLayouts() const {
   for (int i = 0; i < parameter_count(); ++i) {
     TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
         parameter_shape(i),
-        [this, &result](const Shape& subshape, const ShapeIndex& index) {
+        [this, &result](const Shape& subshape,
+                        const ShapeIndex& index) -> absl::Status {
           if (subshape.IsTuple()) {
             return OkStatus();
           }
@@ -93,7 +94,8 @@ absl::StatusOr<std::vector<Layout>> ComputationLayout::FlattenedResultLayouts()
   std::vector<Layout> result;
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
       result_shape(),
-      [this, &result](const Shape& subshape, const ShapeIndex& index) {
+      [this, &result](const Shape& subshape,
+                      const ShapeIndex& index) -> absl::Status {
         if (subshape.IsTuple()) {
           return OkStatus();
         }
diff --git a/xla/service/dynamic_dimension_inference.cc b/xla/service/dynamic_dimension_inference.cc
index 5edb67de0f422..ee74ae367d93f 100644
--- a/xla/service/dynamic_dimension_inference.cc
+++ b/xla/service/dynamic_dimension_inference.cc
@@ -485,8 +485,10 @@ Status DynamicDimensionInferenceVisitor::HandleCustomCall(HloInstruction* hlo) {
     TF_RETURN_IF_ERROR(custom_call_handler_(hlo, parent_));
   } else {
     TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension(
-        hlo, [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-                 int64_t operand_index, HloInstruction* dynamic_size) {
+        hlo,
+        [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
+            int64_t operand_index,
+            HloInstruction* dynamic_size) -> absl::Status {
           // Resize custom call should propagate dynamic batch (0) and channel
           // (3) dimensions.
           if (hlo->custom_call_target() == "SliceToDynamic" ||
@@ -565,8 +567,9 @@ Status DynamicDimensionInferenceVisitor::HandlePad(HloInstruction* hlo) {
     return OkStatus();
   }
   return ForEachOperandDynamicDimension(
-      hlo, [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-               int64_t operand_index, HloInstruction* dynamic_size) {
+      hlo,
+      [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         if (operand_index != 0) {
           return Unimplemented(
               "Dynamic dimension on padding value is not supported");
@@ -803,8 +806,9 @@ Status DynamicDimensionInferenceVisitor::HandleConvolution(
     return OkStatus();
   }
   return ForEachOperandDynamicDimension(
-      hlo, [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-               int64_t operand_index, HloInstruction* dynamic_size) {
+      hlo,
+      [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         HloInstruction* conv = hlo;
         const ConvolutionDimensionNumbers& dimension_numbers =
             conv->convolution_dimension_numbers();
@@ -2120,7 +2124,8 @@ Status DynamicDimensionInferenceVisitor::HandleScatter(HloInstruction* hlo) {
   return ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex dynamic_index, int64_t dimension,
-          int64_t operand_index, HloInstruction* operand_dynamic_size) {
+          int64_t operand_index,
+          HloInstruction* operand_dynamic_size) -> absl::Status {
         if (operand_index == 0) {
           SetDynamicSize(hlo, {}, dimension, operand_dynamic_size);
           return OkStatus();
diff --git a/xla/util.h b/xla/util.h
index c531ff6e97e5f..f726437616244 100644
--- a/xla/util.h
+++ b/xla/util.h
@@ -224,7 +224,6 @@ Status AppendStatus(Status prior, absl::string_view context);
   }
 
 DEFINE_XLA_ERROR_WITH_STRFORMAT_WITH_BACKTRACE(InvalidArgument);
-DEFINE_XLA_ERROR_WITH_STRFORMAT_WITH_BACKTRACE(Unimplemented);
 
 #undef DEFINE_XLA_ERROR_WITH_STRFORMAT_WITH_BACKTRACE
 // The following three macros define a common set of code for creating
@@ -297,6 +296,7 @@ XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(Internal);
 XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(NotFound);
 XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(ResourceExhausted);
 XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(Unavailable);
+XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(Unimplemented);
 XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(Unknown);
 
 #undef XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE