Enable passing down dynamic dimensions from torch to XLA

setup.py

@@ -244,9 +244,10 @@ def bazel_build(self, ext):
 
     bazel_argv = [
         'bazel', 'build', ext.bazel_target,
-        f"--symlink_prefix={os.path.join(self.build_temp, 'bazel-')}",
-        '\n'.join(['--cxxopt=%s' % opt for opt in extra_compile_args])
+        f"--symlink_prefix={os.path.join(self.build_temp, 'bazel-')}"
     ]
+    for opt in extra_compile_args:
+      bazel_argv.append("--cxxopt={}".format(opt))
 
     # Debug build.
     if DEBUG:

torch_xla/csrc/elementwise.cpp

@@ -5,6 +5,7 @@
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/random.h"
 #include "torch_xla/csrc/runtime/debug_macros.h"
+#include "torch_xla/csrc/runtime/sys_util.h"
 #include "torch_xla/csrc/shape_helper.h"
 #include "torch_xla/csrc/tensor_util.h"
 #include "torch_xla/csrc/xla_lower_util.h"
@@ -14,6 +15,9 @@
 namespace torch_xla {
 namespace {
 
+static const bool experimental_unbounded_dynamism =
+    runtime::sys_util::GetEnvBool("EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM", false);
+
 xla::XlaOp Between(xla::XlaOp input, const at::Scalar& min_val,
                    const at::Scalar& max_val) {
   const xla::Shape& shape = ShapeHelper::ShapeOfXlaOp(input);
@@ -66,8 +70,16 @@ xla::XlaOp BuildThreshold(xla::XlaOp input, xla::XlaOp output,
 
 xla::XlaOp BuildRelu(xla::XlaOp input) {
   const xla::Shape& input_shape = ShapeHelper::ShapeOfXlaOp(input);
-  return xla::Max(input, XlaHelpers::ScalarValue<float>(
-                             0, input_shape.element_type(), input.builder()));
+  xla::XlaOp scalar = XlaHelpers::ScalarValue<float>(
+      0, input_shape.element_type(), input.builder());
+  if (experimental_unbounded_dynamism) {
+    // xla::Max doesn't do implicit broadcasting for unbounded dynamism now.
+    // TODO(lsy323): Remove this branch once the support is added in XLA.
+    auto promoted = XlaHelpers::Promote(input, scalar);
+    return xla::Max(promoted.first, promoted.second);
+  } else {
+    return xla::Max(input, scalar);
+  }
 }
 
 xla::XlaOp BuildHardshrink(xla::XlaOp input, xla::XlaOp lambda) {

torch_xla/csrc/helpers.cpp

@@ -21,6 +21,12 @@
 namespace torch_xla {
 namespace {
 
+static const bool experimental_unbounded_dynamism =
+    runtime::sys_util::GetEnvBool("EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM", false);
+
+// TODO(lsy323): Get reserved number for unbounded dim after it's added in XLA.
+static constexpr int64_t kUnboundedSize = std::numeric_limits<int64_t>::min();
+
 xla::XlaOp ConvertBinaryOpResult(xla::XlaOp op1, xla::XlaOp op2,
                                  xla::XlaOp result) {
   xla::PrimitiveType type1 = XlaHelpers::TypeOfXlaOp(op1);
@@ -63,6 +69,9 @@ xla::XlaOp XlaHelpers::BroadcastDimensions(xla::XlaOp input,
   std::vector<int64_t> bcast_sizes = SizesOfXlaOp(input);
   for (size_t i = 0; i < dimensions.size(); ++i) {
     bcast_sizes.at(dimensions[i]) = sizes[i];
+    if (experimental_unbounded_dynamism) {
+      XLA_CHECK(sizes[i] != kUnboundedSize);
+    }
   }
   return xla::BroadcastInDim(input, bcast_sizes,
                              GetAllDimensions(bcast_sizes.size()));
@@ -322,6 +331,57 @@ xla::XlaOp XlaHelpers::DynamicReshapeAs(xla::XlaOp input,
              : xla::Reshape(input, shape.dimensions());
 }
 
+bool XlaHelpers::IsUnboundedDynamic(const xla::Shape& shape) {
+  XLA_CHECK(experimental_unbounded_dynamism)
+      << "EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM needs to be turned on.";
+  const absl::Span<const int64_t> dims = shape.dimensions();
+  return std::any_of(dims.begin(), dims.end(),
+                     [](int64_t size) { return size == kUnboundedSize; });
+}
+
+xla::XlaOp XlaHelpers::DynamicUnboundedReshape(
+    xla::XlaOp input, xla::XlaOp aux_input,
+    absl::Span<const int64_t> output_sizes) {
+  XLA_CHECK(experimental_unbounded_dynamism)
+      << "EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM needs to be turned on.";
+  const xla::Shape& aux_input_shape = ShapeHelper::ShapeOfXlaOp(aux_input);
+  XLA_CHECK(output_sizes.size() == aux_input_shape.rank())
+      << "XlaHelpers::DynamicUnboundedReshape constrainled failed!";
+  std::vector<xla::XlaOp> get_dim_ops;
+  std::vector<xla::XlaOp> reshaped_ops;
+  bool all_static = true;
+  std::vector<bool> output_dynamic(output_sizes.size(), false);
+
+  for (int i = 0; i < output_sizes.size(); i++) {
+    if (output_sizes[i] == kUnboundedSize) {
+      output_dynamic[i] = true;
+      get_dim_ops.push_back(xla::GetDimensionSize(aux_input, i));
+      all_static = false;
+    } else {
+      get_dim_ops.push_back(XlaHelpers::ScalarValue<int32_t>(
+          output_sizes[i], aux_input.builder()));
+    }
+  }
+
+  if (all_static) {
+    return xla::Reshape(input, output_sizes);
+  }
+
+  // Create the reshape from scalar to 1-D vector
+  for (auto get_dim_op : get_dim_ops) {
+    reshaped_ops.push_back(xla::Reshape(get_dim_op, {1}));
+  }
+
+  // Create Concatenate op
+  auto concat_op = xla::ConcatInDim(input.builder(), reshaped_ops, {0});
+  return xla::CustomCall(
+      aux_input.builder(), "stablehlo.dynamic_reshape", {input, concat_op},
+      xla::ShapeUtil::MakeShape(aux_input_shape.element_type(), output_sizes,
+                                output_dynamic));
+
+  return input;
+}
+
 bool XlaHelpers::SameStaticDimensions(const xla::Shape& shape1,
                                       const xla::Shape& shape2) {
   return shape1.is_static() && shape2.is_static() &&
@@ -485,6 +545,11 @@ xla::Shape XlaHelpers::GetPromotedBinaryOpShape(const xla::Shape& shape1,
             runtime::util::ToVector<int64_t>(shape1.dimensions()),
             runtime::util::ToVector<int64_t>(shape2.dimensions())));
   }
+  if (experimental_unbounded_dynamism) {
+    XLA_CHECK(!XlaHelpers::IsUnboundedDynamic(shape1) &&
+              !XlaHelpers::IsUnboundedDynamic(shape2))
+        << "Unreachable for unbounded dynamic code\n";
+  }
   return GetPromotedDynamicShape(shape1, shape2);
 }
 

torch_xla/csrc/helpers.h

@@ -158,6 +158,12 @@ class XlaHelpers {
   static xla::XlaOp DynamicReshape(xla::XlaOp input,
                                    absl::Span<const int64_t> output_sizes);
 
+  static bool IsUnboundedDynamic(const xla::Shape& shape);
+
+  static xla::XlaOp DynamicUnboundedReshape(
+      xla::XlaOp input, xla::XlaOp aux_input,
+      absl::Span<const int64_t> output_sizes);
+
   static xla::XlaOp DynamicReshapeAs(xla::XlaOp input, const xla::Shape& shape);
 
   static bool SameStaticDimensions(const xla::Shape& shape1,

torch_xla/csrc/init_python_bindings.cpp

@@ -1968,6 +1968,12 @@ void InitXlaModuleBindings(py::module m) {
           return handles;
         });
 
+  m.def("_xla_mark_dynamic", [](const at::Tensor& input, uint32_t dim) {
+    TORCH_LAZY_COUNTER("XlaMarkDynamic", 1);
+    XLATensorPtr xtensor = bridge::GetXlaTensor(input);
+    xtensor->MarkDynamicDimension(dim);
+  });
+
   // -------------Dynamo Integration API Start-------------------------
   /*
    * Return tensor ids and at::tensors for all DeviceData nodes that is needed

torch_xla/csrc/ir.cpp

@@ -174,6 +174,10 @@ xla::Shape XlaNode::GetOpShape(
 std::string XlaNode::ToString() const {
   std::stringstream ss;
   ss << torch::lazy::Node::ToString() << ", xla_shape=" << xla_shape_;
+  ss << ", dynamic_dims: ";
+  for (const auto dim : dynamic_dims_) {
+    ss << dim;
+  }
   return ss.str();
 }
 

torch_xla/csrc/ir.h

@@ -9,9 +9,9 @@
 #include <functional>
 #include <iostream>
 #include <memory>
-#include <set>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -138,6 +138,15 @@ class XlaNode : public torch::lazy::Node {
 
   std::string ToString() const override;
 
+  void MarkDynamicDimension(uint32_t dim) { dynamic_dims_.insert(dim); }
+
+  const std::unordered_set<uint32_t>& dynamic_dims() const {
+    return dynamic_dims_;
+  }
+
+ protected:
+  std::unordered_set<uint32_t> dynamic_dims_;
+
  private:
   xla::Shape GetOpShape(const std::function<xla::Shape()>& shape_fn) const;
 

torch_xla/csrc/lowering_context.cpp

@@ -93,19 +93,31 @@ LoweringContext::LoweringContext(
   }
 }
 
+// TODO(lsy323): Get reserved number for unbounded dim after it's added in XLA.
+static constexpr int64_t kUnboundedSize = std::numeric_limits<int64_t>::min();
+
 xla::XlaOp LoweringContext::GetParameter(
-    const std::shared_ptr<torch::lazy::BackendData>& data) {
+    const std::shared_ptr<torch::lazy::BackendData>& data,
+    const std::unordered_set<uint32_t>& unbounded_dynamic_dims) {
   torch::lazy::BackendData::Handle handle = data->GetHandle();
   auto it = parameters_map_.find(handle);
   if (it == parameters_map_.end()) {
-    xla::XlaOp param = xla::Parameter(
-        builder(), parameters_.size(),
+    xla::Shape shape =
         std::dynamic_pointer_cast<runtime::ComputationClient::Data>(data)
-            ->shape(),
-        absl::StrCat("p", parameters_.size()));
+            ->shape();
+    for (const int dim : unbounded_dynamic_dims) {
+      shape.set_dynamic_dimension(dim, true);
+      shape.set_dimensions(dim, kUnboundedSize);
+    }
+    xla::XlaOp param = xla::Parameter(builder(), parameters_.size(), shape,
+                                      absl::StrCat("p", parameters_.size()));
     it = parameters_map_.emplace(handle, Parameter{param, parameters_.size()})
              .first;
     parameters_.push_back(data);
+  } else {
+    XLA_CHECK(unbounded_dynamic_dims.empty())
+        << "The unbounded dynamic dims can only be set when Parameter is "
+           "created.";
   }
   parameter_sequence_.push_back(it->second.index);
   return it->second.param;
@@ -170,6 +182,22 @@ XlaOpVector LoweringContext::LowerNode(const torch::lazy::Node* node) {
 
     const XlaNode* casted = dynamic_cast<const XlaNode*>(node);
     result_ops = casted->Lower(this);
+    if (!casted->dynamic_dims().empty()) {
+      xla::internal::XlaBuilderFriend builder_friend;
+      auto* inst = builder_friend.GetInstruction(result_ops[0]);
+      auto* mutable_dynamic =
+          inst->mutable_shape()->mutable_is_dynamic_dimension();
+      if (mutable_dynamic->empty()) {
+        for (int i = 0; i < inst->dimensions_size(); i++) {
+          mutable_dynamic->Add(false);
+        }
+      }
+      auto* mutable_dims = inst->mutable_shape()->mutable_dimensions();
+      for (const auto dim : casted->dynamic_dims()) {
+        mutable_dynamic->Set(dim, true);
+        mutable_dims->Set(dim, kUnboundedSize);
+      }
+    }
   } catch (const std::exception& ex) {
     ReportBuilderError(node, ex.what());
   }

torch_xla/csrc/lowering_context.h

@@ -37,7 +37,8 @@ class LoweringContext : public torch::lazy::LoweringContext {
   // returned. Otherwise a new one will be created, associated with the tensor
   // held in data.
   xla::XlaOp GetParameter(
-      const std::shared_ptr<torch::lazy::BackendData>& data);
+      const std::shared_ptr<torch::lazy::BackendData>& data,
+      const std::unordered_set<uint32_t>& dynamic_dims = {});
 
   // Retrieves the vector holding all the tensors associated with the parameter
   // instructions which have been created.

torch_xla/csrc/ops/device_data.cpp

@@ -36,7 +36,7 @@ torch::lazy::NodePtr DeviceData::Clone(torch::lazy::OpList operands) const {
 }
 
 XlaOpVector DeviceData::Lower(LoweringContext* loctx) const {
-  return ReturnOp(loctx->GetParameter(data_), loctx);
+  return ReturnOp(loctx->GetParameter(data_, dynamic_dims_), loctx);
 }
 
 DeviceData* DeviceData::Cast(const torch::lazy::Node* node) {

torch_xla/csrc/reduction.cpp

@@ -31,6 +31,9 @@ struct SummationResult {
   xla::XlaOp result;
 };
 
+static const bool experimental_unbounded_dynamism =
+    runtime::sys_util::GetEnvBool("EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM", false);
+
 ReductionInfo GetReductionInfo(xla::XlaOp input, const xla::Shape& shape,
                                absl::Span<const int64_t> dimensions,
                                bool keep_reduced_dimensions) {
@@ -81,7 +84,15 @@ xla::XlaOp GetScaleValue(xla::XlaOp input, xla::XlaOp count,
   xla::XlaOp scale = xla::Select(xla::Ne(count, zero),
                                  one / xla::ConvertElementType(count, type),
                                  xla::NanValue(input.builder(), type));
-  return input * scale;
+
+  if (experimental_unbounded_dynamism) {
+    // XLA Multiply doesn't do implicit broadcasting for unbounded dynamism now.
+    // TODO(lsy323): Remove this branch once the support is added in XLA.
+    auto promoted = XlaHelpers::Promote(input, scale);
+    return promoted.first * promoted.second;
+  } else {
+    return input * scale;
+  }
 }
 
 xla::XlaOp AverageValue(xla::XlaOp input, xla::XlaOp reduced) {
@@ -109,8 +120,15 @@ SummationResult CreateSummation(xla::XlaOp input,
         result.result, result.rinfo.element_count.size, shape.element_type());
   }
   if (keep_reduced_dimensions) {
-    result.result =
-        XlaHelpers::DynamicReshape(result.result, result.rinfo.new_dimensions);
+    if (experimental_unbounded_dynamism) {
+      // TODO(lsy323): Use XLA DynamicReshape once unbounded dynamism support is
+      // added.
+      result.result = XlaHelpers::DynamicUnboundedReshape(
+          result.result, input, result.rinfo.new_dimensions);
+    } else {
+      result.result = XlaHelpers::DynamicReshape(result.result,
+                                                 result.rinfo.new_dimensions);
+    }
   }
   return result;
 }

torch_xla/csrc/tensor.cpp

@@ -893,4 +893,9 @@ int64_t XLATensor::GetHandle() const {
   }
 }
 
+void XLATensor::MarkDynamicDimension(uint32_t dim) {
+  auto* xla_node = dynamic_cast<XlaNode*>(GetIrValue().node.get());
+  xla_node->MarkDynamicDimension(dim);
+}
+
 }  // namespace torch_xla

torch_xla/csrc/tensor.h

@@ -201,6 +201,7 @@ class XLATensor : public torch::lazy::LazyTensor {
   // Set logical_element_type which is visible to upstream PyTorch.
   void SetScalarType(c10::optional<at::ScalarType> logical_element_type);
 
+  void MarkDynamicDimension(uint32_t dim);
   // We don't use the upstream shape to provide xla::shape.
   runtime::util::MaybeRef<xla::Shape> shape() const;
 

torch_xla/stablehlo.py

@@ -212,6 +212,16 @@ def call_function(self, target, args: Tuple, kwargs: Dict) -> Any:
       new_kwargs['device'] = self._device
     return super().call_function(target, args, new_kwargs)
 
+  def run_node(self, n) -> Any:
+    if n.op == 'placeholder':
+      fake_t = n.meta['val']
+      res = super().run_node(n)
+      for i, x in enumerate(fake_t.shape):
+        if not isinstance(x, int):
+          torch_xla._XLAC._xla_mark_dynamic(res, i)
+      return res
+    return super().run_node(n)
+
 
 def _extract_input_args(exported_model, options):
   if options.override_tracing_arguments is not None: