torch_xla: Use new macros for throwing exceptions (part 1). (#9593)

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_xla_bridge.cpp

@@ -170,7 +170,9 @@ torch_xla::XLATensorPtr GetXlaTensorOrCreateForWrappedNumber(
       (tensor.dim() == 0 && tensor.numel() == 1)) {
     return torch_xla::bridge::GetOrCreateXlaTensor(tensor, device);
   } else {
-    return GetValueOrThrow(torch_xla::bridge::GetXlaTensor(tensor));
+    XLA_ASSIGN_OR_THROW(XLATensorPtr xla_tensor,
+                        torch_xla::bridge::GetXlaTensor(tensor));
+    return xla_tensor;
   }
 }
 
@@ -186,9 +188,13 @@ XLATensorPtr GetOrCreateXlaTensor(const at::Tensor& tensor,
   }
 
   auto xtensor = GetXlaTensor(tensor);
-  return xtensor.ok()
-             ? xtensor.value()
-             : GetValueOrThrow(XLATensor::Create(inner_tensor, device));
+  if (xtensor.ok()) {
+    return xtensor.value();
+  }
+
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_tensor,
+                      XLATensor::Create(inner_tensor, device));
+  return xla_tensor;
 }
 
 XLATensorPtr GetOrCreateXlaTensor(const std::optional<at::Tensor>& tensor,
@@ -479,8 +485,8 @@ at::Tensor CreateXlaTensor(
     at::Tensor tensor,
     const std::optional<torch::lazy::BackendDevice>& device) {
   if (tensor.defined() && device) {
-    XLATensorPtr xla_tensor =
-        GetValueOrThrow(XLATensor::Create(std::move(tensor), *device));
+    XLA_ASSIGN_OR_THROW(XLATensorPtr xla_tensor,
+                        XLATensor::Create(std::move(tensor), *device));
     tensor = AtenFromXlaTensor(xla_tensor);
   }
   return tensor;

torch_xla/csrc/debug_util.cpp

@@ -450,9 +450,9 @@ void DebugUtil::post_compilation_analysis(
   // This can be used to verify the hash of the underlying computation proto.
   // Note that for UserComputation computations, the protobuf is factored in
   // the graph hash.
-  std::string serialized_computation =
-      GetValueOrThrow(runtime::util::GetDeterministicSerializedModuleProto(
-          computation->computation().proto()));
+  XLA_ASSIGN_OR_THROW(std::string serialized_computation,
+                      runtime::util::GetDeterministicSerializedModuleProto(
+                          computation->computation().proto()));
   ss << "\n"
      << "Computation hash: "
      << torch::lazy::HashToString(torch::lazy::Hash(serialized_computation))

torch_xla/csrc/dl_convertor.cpp

@@ -141,10 +141,10 @@ DLManagedTensor* toDLPack(const at::Tensor& input) {
   DLTensor& dt = pack->tensor.dl_tensor;
   {
     // AcquireExternalReference may block
-    pack->external_reference =
-        GetValueOrThrow(pjrt_buffer->AcquireExternalReference());
+    XLA_ASSIGN_OR_THROW(pack->external_reference,
+                        pjrt_buffer->AcquireExternalReference());
     xla::PjRtFuture<> future = pjrt_buffer->GetReadyFuture();
-    OkOrThrow(future.Await());
+    XLA_THROW_IF_ERROR(future.Await());
   }
   pack->buffer_reference = pjrt_buffer;
 
@@ -329,8 +329,9 @@ at::Tensor fromDLPack(DLManagedTensor* dlmt) {
   if (dlmt->deleter) {
     on_delete_callback = [dlmt]() { dlmt->deleter(dlmt); };
   }
-  std::unique_ptr<xla::PjRtBuffer> pjrt_buffer =
-      GetValueOrThrow(device->client()->CreateViewOfDeviceBuffer(
+  XLA_ASSIGN_OR_THROW(
+      std::unique_ptr<xla::PjRtBuffer> pjrt_buffer,
+      device->client()->CreateViewOfDeviceBuffer(
           static_cast<char*>(dlmt->dl_tensor.data) +
               dlmt->dl_tensor.byte_offset,
           shape, *device->default_memory_space(), on_delete_callback));

torch_xla/csrc/helpers.cpp

@@ -41,7 +41,8 @@ xla::XlaComputation CreateComputation(
       xla::Parameter(&builder, 0, xla::ShapeUtil::MakeShape(type, {}), "x");
   xla::XlaOp y =
       xla::Parameter(&builder, 1, xla::ShapeUtil::MakeShape(type, {}), "y");
-  return GetValueOrThrow(builder.Build(op(x, y)));
+  XLA_ASSIGN_OR_THROW(xla::XlaComputation computation, builder.Build(op(x, y)));
+  return computation;
 }
 
 xla::XlaComputation CreateMinMaxComputation(const std::string& name,
@@ -66,7 +67,8 @@ xla::XlaComputation CreateMinMaxComputation(const std::string& name,
   xla::XlaOp tie_id = xla::Min(lhs_index, rhs_index);
   arg_max = xla::Select(eq, tie_id, arg_max);
   xla::Tuple(&builder, {max, arg_max});
-  return GetValueOrThrow(builder.Build());
+  XLA_ASSIGN_OR_THROW(xla::XlaComputation min_max_computation, builder.Build());
+  return min_max_computation;
 }
 
 }  // namespace
@@ -697,7 +699,8 @@ std::vector<int64_t> XlaHelpers::getBroadcastDimensions(xla::XlaOp op1,
 xla::Shape XlaHelpers::GetPromotedBinaryOpShape(const xla::Shape& shape1,
                                                 const xla::Shape& shape2) {
   if (!shape1.is_dynamic() && !shape2.is_dynamic()) {
-    auto promoted_shape = GetValueOrThrow(GetPromotedShape(shape1, shape2));
+    XLA_ASSIGN_OR_THROW(xla::Shape promoted_shape,
+                        GetPromotedShape(shape1, shape2));
     return xla::ShapeUtil::MakeShape(
         PromoteType(shape1.element_type(), shape2.element_type()),
         promoted_shape.dimensions());
@@ -776,7 +779,7 @@ std::pair<xla::XlaOp, xla::XlaOp> XlaHelpers::PromoteShapes(xla::XlaOp op1,
   const xla::Shape& shape1 = ShapeHelper::ShapeOfXlaOp(op1);
   const xla::Shape& shape2 = ShapeHelper::ShapeOfXlaOp(op2);
 
-  xla::Shape shape = GetValueOrThrow(GetPromotedShape(shape1, shape2));
+  XLA_ASSIGN_OR_THROW(xla::Shape shape, GetPromotedShape(shape1, shape2));
   if (shape1.is_unbounded_dynamic() || shape2.is_unbounded_dynamic()) {
     return ImplicitBroadcastWithUnboundedDynamicShapes(op1, op2, shape);
   }

torch_xla/csrc/lowering_context.cpp

@@ -325,7 +325,7 @@ void LoweringContext::AddParameter(const torch::lazy::Output& output,
 }
 
 torch::lazy::ComputationPtr LoweringContext::Build() {
-  xla::XlaComputation xla_computation = GetValueOrThrow(BuildXla());
+  XLA_ASSIGN_OR_THROW(xla::XlaComputation xla_computation, BuildXla());
   return std::make_shared<runtime::ComputationClient::Computation>(
       builder_.name(), std::move(xla_computation), device_);
 }

torch_xla/csrc/tensor.cpp

@@ -518,8 +518,8 @@ at::Tensor XLATensor::ToTensor(bool detached) {
     XLAGraphExecutor::Get()->DeviceBarrier(GetDevice());
     // The GetXlaData() call will trigger an ApplyPendingGraph() if an IR
     // XlaNode is available on the tensor.
-    std::vector<at::Tensor> tensors =
-        GetValueOrThrow(XlaDataToTensors({GetXlaData()}, {dtype()}));
+    XLA_ASSIGN_OR_THROW(std::vector<at::Tensor> tensors,
+                        XlaDataToTensors({GetXlaData()}, {dtype()}));
     tensor = std::move(tensors.front());
     if (!detached) {
       SetTensorData(tensor);
@@ -627,7 +627,9 @@ std::vector<XLATensorPtr> XLATensor::MakeOutputTensors(
 XLATensorPtr XLATensor::CopyTensorToDevice(
     const torch::lazy::BackendDevice& device) {
   // TODO: This can be optimized via proper XRT/XLA computation.
-  return GetValueOrThrow(Create(ToTensor(/*detached=*/true), device));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr result,
+                      Create(ToTensor(/*detached=*/true), device));
+  return result;
 }
 
 torch::lazy::Value XLATensor::MaybeCastIrValue(

torch_xla/csrc/xla_backend_impl.cpp

@@ -94,7 +94,9 @@ class XlaBackendImpl : public torch::lazy::BackendImplInterface {
       const torch::lazy::BackendDataPtr data,
       std::optional<at::ScalarType> logical_scalar_type) const override {
     // TODO(JackCaoG): handle the logical_scalar_type == nullptr case
-    return GetValueOrThrow(XlaDataToTensors({data}, {*logical_scalar_type}))[0];
+    XLA_ASSIGN_OR_THROW(std::vector<at::Tensor> tensors,
+                        XlaDataToTensors({data}, {*logical_scalar_type}));
+    return tensors[0];
   }
 
   std::unique_ptr<torch::lazy::LoweringContext> CreateLoweringContext(
@@ -163,7 +165,8 @@ class XlaBackendImpl : public torch::lazy::BackendImplInterface {
       torch::lazy::ComputationPtr computation,
       c10::ArrayRef<torch::lazy::BackendDataPtr> arguments,
       const torch::lazy::BackendDevice& device) const override {
-    std::vector<runtime::ComputationClient::DataPtr> results = GetValueOrThrow(
+    XLA_ASSIGN_OR_THROW(
+        std::vector<runtime::ComputationClient::DataPtr> results,
         runtime::GetComputationClientOrDie()->ExecuteComputation(
             *std::dynamic_pointer_cast<runtime::ComputationClient::Computation>(
                 computation),

torch_xla/csrc/xla_manual_registration.cpp

@@ -38,8 +38,8 @@ at::Tensor nms_kernel(const at::Tensor& boxes, const at::Tensor& scores,
   XLA_CHECK_EQ(boxes.size(0), scores.size(0))
       << "nms(): boxes and scores should have the same size for dimension 0.";
 
-  XLATensorPtr xla_boxes = GetValueOrThrow(bridge::GetXlaTensor(boxes));
-  XLATensorPtr xla_scores = GetValueOrThrow(bridge::GetXlaTensor(scores));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_boxes, bridge::GetXlaTensor(boxes));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_scores, bridge::GetXlaTensor(scores));
   return bridge::AtenFromXlaTensor(
       tensor_methods::nms(xla_boxes, xla_scores, iou_threshold),
       /*skip_functionalization=*/true);

torch_xla/csrc/xla_op_builder.cpp

@@ -690,8 +690,8 @@ xla::XlaOp ConcatInDim(const BuilderPtr& builder,
 xla::XlaOp Convert(const BuilderPtr& builder,
                    const std::vector<OpPtr>& operands, py::dict args) {
   std::string type = args["to_type"].cast<std::string>();
-  xla::PrimitiveType xla_type =
-      GetValueOrThrow(xla::primitive_util::StringToPrimitiveType(type));
+  XLA_ASSIGN_OR_THROW(xla::PrimitiveType xla_type,
+                      xla::primitive_util::StringToPrimitiveType(type));
   return MaybeConvertTo(operands.at(0)->op, xla_type);
 }
 
@@ -717,8 +717,8 @@ xla::XlaOp SetDimensionSize(const BuilderPtr& builder,
 xla::XlaOp BitcastConvert(const BuilderPtr& builder,
                           const std::vector<OpPtr>& operands, py::dict args) {
   std::string type = args["to_type"].cast<std::string>();
-  xla::PrimitiveType xla_type =
-      GetValueOrThrow(xla::primitive_util::StringToPrimitiveType(type));
+  XLA_ASSIGN_OR_THROW(xla::PrimitiveType xla_type,
+                      xla::primitive_util::StringToPrimitiveType(type));
   return xla::BitcastConvertType(operands.at(0)->op, xla_type);
 }
 
@@ -873,8 +873,8 @@ xla::Shape PyShapeToShape(py::object shape) {
   std::string type = py_shape["type"].cast<std::string>();
   std::vector<int64_t> dimensions =
       GetTupleVector<int64_t>(py_shape["sizes"].cast<py::tuple>());
-  xla::PrimitiveType xla_type =
-      GetValueOrThrow(xla::primitive_util::StringToPrimitiveType(type));
+  XLA_ASSIGN_OR_THROW(xla::PrimitiveType xla_type,
+                      xla::primitive_util::StringToPrimitiveType(type));
   if (py_shape.contains("dynamic_dimensions")) {
     std::vector<bool> dynamic_dimensions =
         GetTupleVector<bool>(py_shape["dynamic_dimensions"]);

torch_xla/csrc/xla_sharding_util.cpp

@@ -767,7 +767,7 @@ void ShardingUtil::XlaMarkSharding(const at::Tensor& input,
       << "Please enable SPMD via `torch_xla.runtime.use_spmd()`";
   XLA_CHECK(sharding.type() != xla::OpSharding::UNKNOWN)
       << "Can't explicilty annotate with UNKNOWN sharding type.";
-  XLATensorPtr xtensor = GetValueOrThrow(bridge::GetXlaTensor(input));
+  XLA_ASSIGN_OR_THROW(XLATensorPtr xtensor, bridge::GetXlaTensor(input));
 
   // For Non DeviceData IR values, we directly attach the sharding spec to the
   // xtensor.