feat: more JLL changes for #788 (#790)

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -78,6 +78,7 @@
 #include "shardy/integrations/c/attributes.h"
 #include "xla/pjrt/mlir_to_hlo.h"
 #include "xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.h"
+#include "xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h"
 
 // IFRT
 #include "xla/python/ifrt/array.h"
@@ -1956,4 +1957,31 @@ hloShardingFromTensorShardingAttr(mlir::sdy::TensorShardingAttr attr,
       xla::sdy::convertToHloSharding(attr, get_mesh_attr, manual_axes));
 }
 
+extern "C" mlir::sdy::TensorShardingAttr
+hloShardingToTensorShardingAttr(const xla::HloSharding *hloSharding,
+                                mlir::sdy::MeshAttr meshAttr, int64_t rank,
+                                const bool *isClosed, const int64_t *priority) {
+  const SmallDenseMap<int64_t, StringRef> deviceIdToMaximalMeshName =
+      SmallDenseMap<int64_t, StringRef>();
+  mlir::sdy::TensorShardingAttr tensorShardingAttr =
+      xla::sdy::convertToSdySharding(*hloSharding, meshAttr,
+                                     deviceIdToMaximalMeshName, rank,
+                                     /*openDims=*/true);
+
+  for (int64_t i = 0; i < rank; i++) {
+    auto oldDimSharding = tensorShardingAttr.getDimSharding(i);
+
+    std::optional<int64_t> dimPriority;
+    if (priority[i] > 0)
+      dimPriority = priority[i];
+
+    tensorShardingAttr = tensorShardingAttr.replaceDimSharding(
+        i, mlir::sdy::DimensionShardingAttr::get(oldDimSharding.getContext(),
+                                                 oldDimSharding.getAxes(),
+                                                 isClosed[i], dimPriority));
+  }
+
+  return tensorShardingAttr;
+}
+
 #pragma endregion

deps/ReactantExtra/BUILD

@@ -492,6 +492,7 @@ cc_library(
 "-Wl,-exported_symbol,_ClientGetAddressableDevices",
 "-Wl,-exported_symbol,_hloShardingFromTensorShardingAttr",
 "-Wl,-exported_symbol,_op_sharding_*",
+"-Wl,-exported_symbol,_hloShardingToTensorShardingAttr",
     ]}),
     deps = [
                 "@enzyme//:EnzymeMLIR",
@@ -541,6 +542,7 @@ cc_library(
         "@xla//xla/pjrt/distributed:client",
         "@xla//xla/pjrt/distributed:service",
         "@xla//xla/service/spmd/shardy/stablehlo_round_trip:export_shardings",
+        "@xla//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import",
 
         "@xla//xla:xla_proto_cc",
         "@xla//xla:xla_proto_cc_impl",