torch_xla: Use new macros for throwing exceptions (part 2). (#9594)

Follow-up: #9588 and #9580
Target: `torch_xla/csrc` directory

In summary, this PR:

- Replaces all calls to `OkOrThrow()` and `GetValueOrThrow()` (that
throws an exception without source location information of the
*"throw-site"*) with the macros `XLA_THROW_IF_ERROR()` and
`XLA_ASSIGN_OR_THROW()`.
- Corresponds to the fine-grained set of PRs that came from breaking
down PR #9580
- Focuses on the `torch_xla/csrc` directory, replacing every use of
those, now deprecated, functions by the newly introduced macros.

_Note: since there were lots of files in `torch_xla/csrc` that needed
update, they were split in multiple parts._

Author: ysiraichi

torch_xla/csrc/aten_autograd_ops.cpp

@@ -34,8 +34,8 @@ torch::Tensor EinsumAutogradFunction::forward(
   }
   ctx->save_for_backward(vars);
 
-  std::vector<XLATensorPtr> xla_tensors =
-      GetValueOrThrow(bridge::GetXlaTensors(tensors));
+  XLA_ASSIGN_OR_THROW(std::vector<absl_nonnull XLATensorPtr> xla_tensors,
+                      bridge::GetXlaTensors(tensors));
   XLATensorPtr output = tensor_methods::einsum(eq_str, xla_tensors);
   return bridge::AtenFromXlaTensor(output);
 }
@@ -45,13 +45,12 @@ torch::autograd::variable_list EinsumAutogradFunction::backward(
     torch::autograd::variable_list grad_output) {
   std::string equation = ctx->saved_data["equation"].toString()->string();
   torch::autograd::variable_list tensors = ctx->get_saved_variables();
-  std::vector<XLATensorPtr> xla_tensors =
-      GetValueOrThrow(bridge::GetXlaTensors(tensors));
-
+  XLA_ASSIGN_OR_THROW(std::vector<absl_nonnull XLATensorPtr> xla_tensors,
+                      bridge::GetXlaTensors(tensors));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_grad_output_0,
+                      bridge::GetXlaTensor(grad_output[0]));
   std::tuple<XLATensorPtr, XLATensorPtr> outputs =
-      tensor_methods::einsum_backward(
-          GetValueOrThrow(bridge::GetXlaTensor(grad_output[0])), xla_tensors,
-          equation);
+      tensor_methods::einsum_backward(xla_grad_output_0, xla_tensors, equation);
 
   // For both einsum and max pool, we use "undef" as a placeholder for the
   // non-tensor grad inputs, in this case the equation string.
@@ -193,10 +192,10 @@ torch::Tensor MaxPool3dAutogradFunction::forward(
     return std::get<0>(results);
   }
   ctx->save_for_backward({self});
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_self, bridge::GetXlaTensor(self));
   auto outputs = tensor_methods::max_pool_nd(
-      GetValueOrThrow(bridge::GetXlaTensor(self)), /*spatial_dim_count=*/3,
-      XlaHelpers::I64List(kernel_size), XlaHelpers::I64List(stride),
-      XlaHelpers::I64List(padding), ceil_mode);
+      xla_self, /*spatial_dim_count=*/3, XlaHelpers::I64List(kernel_size),
+      XlaHelpers::I64List(stride), XlaHelpers::I64List(padding), ceil_mode);
   return bridge::AtenFromXlaTensor(std::get<0>(outputs));
 }
 
@@ -221,11 +220,13 @@ torch::autograd::variable_list MaxPool3dAutogradFunction::backward(
                                                          padding, dilation,
                                                          ceil_mode, indices);
   }
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_grad_output_0,
+                      bridge::GetXlaTensor(grad_output[0]));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_self, bridge::GetXlaTensor(self));
   grad = bridge::AtenFromXlaTensor(tensor_methods::max_pool_nd_backward(
-      GetValueOrThrow(bridge::GetXlaTensor(grad_output[0])),
-      GetValueOrThrow(bridge::GetXlaTensor(self)),
-      /*spatial_dim_count=*/3, XlaHelpers::I64List(kernel_size),
-      XlaHelpers::I64List(stride), XlaHelpers::I64List(padding), ceil_mode));
+      xla_grad_output_0, xla_self, /*spatial_dim_count=*/3,
+      XlaHelpers::I64List(kernel_size), XlaHelpers::I64List(stride),
+      XlaHelpers::I64List(padding), ceil_mode));
 
   torch::Tensor undef;
   torch::autograd::variable_list grad_inputs = {grad,  undef, undef,
@@ -238,22 +239,24 @@ torch::Tensor max_pool2d_forward(torch::Tensor self,
                                  torch::IntArrayRef stride,
                                  torch::IntArrayRef padding,
                                  torch::IntArrayRef dilation, bool ceil_mode) {
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_self, bridge::GetXlaTensor(self));
   auto outputs = tensor_methods::max_pool_nd(
-      GetValueOrThrow(bridge::GetXlaTensor(self)), /*spatial_dim_count=*/2,
-      XlaHelpers::I64List(kernel_size), XlaHelpers::I64List(stride),
-      XlaHelpers::I64List(padding), ceil_mode);
+      xla_self, /*spatial_dim_count=*/2, XlaHelpers::I64List(kernel_size),
+      XlaHelpers::I64List(stride), XlaHelpers::I64List(padding), ceil_mode);
   return bridge::AtenFromXlaTensor(std::get<0>(outputs));
 }
 
 torch::Tensor max_pool2d_backward(torch::Tensor grad_output, torch::Tensor self,
                                   torch::IntArrayRef kernel_size,
                                   torch::IntArrayRef stride,
                                   torch::IntArrayRef padding, bool ceil_mode) {
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_grad_output,
+                      bridge::GetXlaTensor(grad_output));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_self, bridge::GetXlaTensor(self));
   auto grad = bridge::AtenFromXlaTensor(tensor_methods::max_pool_nd_backward(
-      GetValueOrThrow(bridge::GetXlaTensor(grad_output)),
-      GetValueOrThrow(bridge::GetXlaTensor(self)),
-      /*spatial_dim_count=*/2, XlaHelpers::I64List(kernel_size),
-      XlaHelpers::I64List(stride), XlaHelpers::I64List(padding), ceil_mode));
+      xla_grad_output, xla_self, /*spatial_dim_count=*/2,
+      XlaHelpers::I64List(kernel_size), XlaHelpers::I64List(stride),
+      XlaHelpers::I64List(padding), ceil_mode));
   return grad;
 }
 

torch_xla/csrc/ir_dump_util.cpp

@@ -264,14 +264,15 @@ std::string DumpUtil::ToHlo(c10::ArrayRef<torch::lazy::Value> values,
   // Annotate HLO sharding selectively in the compuation.
   // This is no-op if an instruction doesn't have any sharding annotation.
   auto is_sharded = ShardingUtil::SetHloSharding(&lowering_ctx);
-  xla::XlaComputation computation = GetValueOrThrow(lowering_ctx.BuildXla());
+  XLA_ASSIGN_OR_THROW(xla::XlaComputation computation, lowering_ctx.BuildXla());
 
   static bool dump_post_optimizations =
       runtime::sys_util::GetEnvBool("XLA_DUMP_POST_OPTIMIZATIONS", false);
   if (dump_post_optimizations) {
+    XLA_ASSIGN_OR_THROW(xla::ProgramShape program_shape,
+                        computation.GetProgramShape());
     xla::Shape shape = MakeShapeWithDeviceLayout(
-        GetValueOrThrow(computation.GetProgramShape()).result(),
-        static_cast<XlaDeviceType>(device.type()));
+        program_shape.result(), static_cast<XlaDeviceType>(device.type()));
     std::vector<runtime::ComputationClient::CompileInstance> instances;
     instances.push_back(
         {std::move(computation), device.toString(),
@@ -286,12 +287,17 @@ std::string DumpUtil::ToHlo(c10::ArrayRef<torch::lazy::Value> values,
   }
 
   switch (mode) {
-    case EmitMode::kHloReadable:
-      return GetValueOrThrow(runtime::util::GetComputationHloText(computation));
-    case EmitMode::kHloProto:
-      return GetValueOrThrow(
-          runtime::util::GetDeterministicSerializedModuleProto(
-              computation.proto()));
+    case EmitMode::kHloReadable: {
+      XLA_ASSIGN_OR_THROW(std::string hlo_text,
+                          runtime::util::GetComputationHloText(computation));
+      return hlo_text;
+    }
+    case EmitMode::kHloProto: {
+      XLA_ASSIGN_OR_THROW(std::string serialized_proto,
+                          runtime::util::GetDeterministicSerializedModuleProto(
+                              computation.proto()));
+      return serialized_proto;
+    }
     case EmitMode::kStableHloReadable:
       return hloToStablehlo(&computation.proto(),
                             /* emit_bytecode = */ false);

torch_xla/csrc/xla_graph_executor.cpp

@@ -497,8 +497,8 @@ std::vector<at::Tensor> XLAGraphExecutor::GetTensors(
       async != nullptr ? async->tensors_data
                        : absl::Span<const torch::lazy::BackendDataPtr>());
 
-  std::vector<xla::Literal> literals =
-      GetValueOrThrow(ReleaseGilAndTransferData(tensors_data));
+  XLA_ASSIGN_OR_THROW(std::vector<xla::Literal> literals,
+                      ReleaseGilAndTransferData(tensors_data));
 
   return FetchTensors(tensors, literals,
                       async != nullptr ? &async->indices : nullptr);
@@ -846,12 +846,12 @@ XLAGraphExecutor::ExecuteComputationWithBarrier(
         // OutputHandler creates sharded data for sharded
         // tensor results. Both sharded and unsharded results should be
         // "Assign"ed to the corresponding data placeholders.
-        std::vector<runtime::ComputationClient::DataPtr> outputs =
-            GetValueOrThrow(
-                runtime::GetComputationClientOrDie()->ExecuteReplicated(
-                    *async->cached_computation->computation,
-                    UnwrapXlaData(async->parameters_data), devices,
-                    execute_options));
+        XLA_ASSIGN_OR_THROW(
+            std::vector<runtime::ComputationClient::DataPtr> outputs,
+            runtime::GetComputationClientOrDie()->ExecuteReplicated(
+                *async->cached_computation->computation,
+                UnwrapXlaData(async->parameters_data), devices,
+                execute_options));
         results = WrapXlaData(outputs);
         TF_VLOG(3) << "Executing Dynamo IR sharded graph hash "
                    << torch::lazy::HashToString(hash) << " on devices "
@@ -913,8 +913,8 @@ std::vector<torch::lazy::BackendDataPtr> XLAGraphExecutor::ExecuteStablehlo(
 
   // Get program output shape.
   // TODO(lsy323): Get shape info from MLIR Module.
-  xla::ProgramShape program_shape =
-      GetValueOrThrow(computation.GetProgramShape());
+  XLA_ASSIGN_OR_THROW(xla::ProgramShape program_shape,
+                      computation.GetProgramShape());
   xla::Shape shape = MakeShapeWithDeviceLayout(
       program_shape.result(), static_cast<XlaDeviceType>(device.type()));
 
@@ -946,8 +946,9 @@ std::vector<torch::lazy::BackendDataPtr> XLAGraphExecutor::ExecuteStablehlo(
     }
   }
 
-  std::vector<runtime::ComputationClient::DataPtr> result_data =
-      GetValueOrThrow(runtime::GetComputationClientOrDie()->ExecuteComputation(
+  XLA_ASSIGN_OR_THROW(
+      std::vector<runtime::ComputationClient::DataPtr> result_data,
+      runtime::GetComputationClientOrDie()->ExecuteComputation(
           *computations[0], UnwrapXlaData(arguments), device.toString()));
 
   return WrapXlaData(result_data);
@@ -1123,12 +1124,12 @@ XLAGraphExecutor::ScheduleSyncTensorsGraph(
         // OutputHandler creates sharded data for sharded
         // tensor results. Both sharded and unsharded results should be
         // "Assign"ed to the corresponding data placeholders.
-        std::vector<runtime::ComputationClient::DataPtr> outputs =
-            GetValueOrThrow(
-                runtime::GetComputationClientOrDie()->ExecuteReplicated(
-                    *async->cached_computation->computation,
-                    UnwrapXlaData(async->parameters_data), devices,
-                    execute_options));
+        XLA_ASSIGN_OR_THROW(
+            std::vector<runtime::ComputationClient::DataPtr> outputs,
+            runtime::GetComputationClientOrDie()->ExecuteReplicated(
+                *async->cached_computation->computation,
+                UnwrapXlaData(async->parameters_data), devices,
+                execute_options));
         results = WrapXlaData(outputs);
         TORCH_LAZY_COUNTER("ExecuteReplicated", 1);
         TF_VLOG(3) << "Executing IR graph hash "
@@ -1139,14 +1140,13 @@ XLAGraphExecutor::ScheduleSyncTensorsGraph(
         TF_VLOG(3) << "Executing IR graph hash "
                    << torch::lazy::HashToString(hash) << " on device "
                    << async->device << " ...";
-        std::vector<runtime::ComputationClient::DataPtr> outputs =
-            GetValueOrThrow(
-                runtime::GetComputationClientOrDie()->ExecuteComputation(
-                    *async->cached_computation->computation,
-                    UnwrapXlaData(async->parameters_data),
-                    async->device.toString(),
-                    {/*explode_tuple=*/true,
-                     /*eager_mode=*/use_eager_mode}));
+        XLA_ASSIGN_OR_THROW(
+            std::vector<runtime::ComputationClient::DataPtr> outputs,
+            runtime::GetComputationClientOrDie()->ExecuteComputation(
+                *async->cached_computation->computation,
+                UnwrapXlaData(async->parameters_data), async->device.toString(),
+                {/*explode_tuple=*/true,
+                 /*eager_mode=*/use_eager_mode}));
         results = WrapXlaData(outputs);
         TORCH_LAZY_COUNTER("ExecuteComputation", 1);
         TF_VLOG(3) << "Executing IR graph hash "
@@ -1416,9 +1416,9 @@ XLAGraphExecutor::CompilationResult XLAGraphExecutor::Compile(
 
   SetBufferDonors(&lowering_ctx, buffer_donor_indices);
 
-  xla::XlaComputation computation = GetValueOrThrow(lowering_ctx.BuildXla());
-  xla::ProgramShape program_shape =
-      GetValueOrThrow(computation.GetProgramShape());
+  XLA_ASSIGN_OR_THROW(xla::XlaComputation computation, lowering_ctx.BuildXla());
+  XLA_ASSIGN_OR_THROW(xla::ProgramShape program_shape,
+                      computation.GetProgramShape());
 
   // TODO(yeounoh) enable wrapping with auto-sharding.
   bool should_wrap_parameter =
@@ -1435,10 +1435,11 @@ XLAGraphExecutor::CompilationResult XLAGraphExecutor::Compile(
       param_shardings = XlaHelpers::ExtractInputShardings(computation);
     }
 
-    computation = GetValueOrThrow(
+    XLA_ASSIGN_OR_THROW(
+        computation,
         XlaHelpers::WrapXlaComputation(computation, program_shape.parameters(),
                                        param_shardings, buffer_donor_indices));
-    program_shape = GetValueOrThrow(computation.GetProgramShape());
+    XLA_ASSIGN_OR_THROW(program_shape, computation.GetProgramShape());
   }
   xla::Shape shape = MakeShapeWithDeviceLayout(
       program_shape.result(), static_cast<XlaDeviceType>(coll.device.type()));

torch_xla/csrc/xla_lower_util.cpp

@@ -62,8 +62,8 @@ ConditionMaskData CreateConditionMaskData(xla::XlaOp condition) {
 
 xla::XlaOp GetPromotedMask(xla::XlaOp mask, const xla::Shape& input_shape) {
   const xla::Shape& mask_shape = ShapeHelper::ShapeOfXlaOp(mask);
-  xla::Shape promoted_mask_shape =
-      GetValueOrThrow(XlaHelpers::GetPromotedShape(mask_shape, input_shape));
+  XLA_ASSIGN_OR_THROW(xla::Shape promoted_mask_shape,
+                      XlaHelpers::GetPromotedShape(mask_shape, input_shape));
   return XlaHelpers::ImplicitBroadcast(mask, mask_shape, promoted_mask_shape);
 }
 
@@ -150,7 +150,9 @@ xla::XlaComputation MakeScatterComputation(
   if (combiner != nullptr) {
     result = combiner(p0, result);
   }
-  return GetValueOrThrow(cb.Build(result));
+  XLA_ASSIGN_OR_THROW(xla::XlaComputation scatter_computation,
+                      cb.Build(result));
+  return scatter_computation;
 }
 
 xla::XlaOp CreateIndexAlongDim(
@@ -543,8 +545,8 @@ std::vector<xla::XlaOp> CreateBroadcastTensors(
   for (const xla::XlaOp operand : operands) {
     const xla::Shape& operand_shape = ShapeHelper::ShapeOfXlaOp(operand);
     operand_shapes.push_back(operand_shape);
-    result_shape = GetValueOrThrow(
-        XlaHelpers::GetPromotedShape(result_shape, operand_shape));
+    XLA_ASSIGN_OR_THROW(result_shape, XlaHelpers::GetPromotedShape(
+                                          result_shape, operand_shape));
   }
   std::vector<xla::XlaOp> result;
   for (size_t i = 0; i < operands.size(); ++i) {
@@ -1366,54 +1368,59 @@ std::vector<xla::XlaOp> BuildBoxSelectionLoop(int64_t num_boxes,
   // 3. The actual IoU threshold matrix.
   init_values[2] = iou_threshold_mask;
 
-  return GetValueOrThrow(xla::WhileLoopHelper(
-      [=](absl::Span<const xla::XlaOp> values, xla::XlaBuilder* builder) {
-        xla::XlaOp box_index = values[0];
-        // Check: current loop counter is within bounds, i.e. has a
-        // corresponding box.
-        return xla::Lt(box_index,
-                       xla::ConstantR0<IndexType>(builder, num_boxes));
-      },
-      [=](absl::Span<const xla::XlaOp> values, xla::XlaBuilder* builder) {
-        const xla::XlaOp ONE = xla::One(builder, XLAIndexType);
-        const xla::XlaOp ZERO = xla::Zero(builder, XLAIndexType);
-
-        xla::XlaOp box_index = values[0];
-        xla::XlaOp state = values[1];
-        xla::XlaOp iou_threshold_mask = values[2];
-
-        // Retrieve the IoU mask row corresponding to this box.
-        xla::XlaOp box_iou_threshold_mask = xla::DynamicSlice(
-            iou_threshold_mask, {box_index, ZERO}, {1, num_boxes});
-
-        // Update the current state with the IoU mask.
-        // Basically, sets to false every box X whose IoU with the current box
-        // is less-than or equal than the given threshold.
-        xla::XlaOp updated_state = xla::And(
-            state,
-            // Update the mask so that if we select this box
-            // (i.e. state[box] == true), we don't de-select it.
-            xla::DynamicUpdateSlice(
-                // Before that, we need to pre-process the mask.
-                //   1. Negate the mask: if this box is selected, we only want
-                //      those that have a low intersection ratio.
-                //   2. Reshape it to: [num_boxes].
-                xla::Reshape(xla::Not(box_iou_threshold_mask), {num_boxes}),
-                xla::ConstantR1<bool>(builder, {true}), {box_index}));
-
-        // Flag: should this box (loop counter) be included in the output?
-        xla::XlaOp should_include = xla::DynamicSlice(state, {box_index}, {1});
-        // Pick the new values of state, depending on whether we should include
-        // this box or not.
-        xla::XlaOp new_state =
-            xla::Select(xla::BroadcastInDim(should_include, {num_boxes}, {0}),
-                        updated_state, state);
-
-        xla::XlaOp next_box_index = box_index + ONE;
-        return std::vector<xla::XlaOp>{next_box_index, new_state,
-                                       iou_threshold_mask};
-      },
-      init_values, "BoxSelectionLoop", builder));
+  XLA_ASSIGN_OR_THROW(
+      std::vector<xla::XlaOp> result,
+      xla::WhileLoopHelper(
+          [=](absl::Span<const xla::XlaOp> values, xla::XlaBuilder* builder) {
+            xla::XlaOp box_index = values[0];
+            // Check: current loop counter is within bounds, i.e. has a
+            // corresponding box.
+            return xla::Lt(box_index,
+                           xla::ConstantR0<IndexType>(builder, num_boxes));
+          },
+          [=](absl::Span<const xla::XlaOp> values, xla::XlaBuilder* builder) {
+            const xla::XlaOp ONE = xla::One(builder, XLAIndexType);
+            const xla::XlaOp ZERO = xla::Zero(builder, XLAIndexType);
+
+            xla::XlaOp box_index = values[0];
+            xla::XlaOp state = values[1];
+            xla::XlaOp iou_threshold_mask = values[2];
+
+            // Retrieve the IoU mask row corresponding to this box.
+            xla::XlaOp box_iou_threshold_mask = xla::DynamicSlice(
+                iou_threshold_mask, {box_index, ZERO}, {1, num_boxes});
+
+            // Update the current state with the IoU mask.
+            // Basically, sets to false every box X whose IoU with the current
+            // box is less-than or equal than the given threshold.
+            xla::XlaOp updated_state = xla::And(
+                state,
+                // Update the mask so that if we select this box
+                // (i.e. state[box] == true), we don't de-select it.
+                xla::DynamicUpdateSlice(
+                    // Before that, we need to pre-process the mask.
+                    //   1. Negate the mask: if this box is selected, we only
+                    //   want
+                    //      those that have a low intersection ratio.
+                    //   2. Reshape it to: [num_boxes].
+                    xla::Reshape(xla::Not(box_iou_threshold_mask), {num_boxes}),
+                    xla::ConstantR1<bool>(builder, {true}), {box_index}));
+
+            // Flag: should this box (loop counter) be included in the output?
+            xla::XlaOp should_include =
+                xla::DynamicSlice(state, {box_index}, {1});
+            // Pick the new values of state, depending on whether we should
+            // include this box or not.
+            xla::XlaOp new_state = xla::Select(
+                xla::BroadcastInDim(should_include, {num_boxes}, {0}),
+                updated_state, state);
+
+            xla::XlaOp next_box_index = box_index + ONE;
+            return std::vector<xla::XlaOp>{next_box_index, new_state,
+                                           iou_threshold_mask};
+          },
+          init_values, "BoxSelectionLoop", builder));
+  return result;
 }
 
 xla::XlaOp BuildNms(xla::XlaOp boxes, xla::XlaOp scores,