Build absl::string_view(data, length) (instead of StringRef::str) exp…

…licitly since the llvm::StringRef to absl::string_view converter is not (always?) available on Android. END_PUBLIC PiperOrigin-RevId: 730637186
openxla · Feb 25, 2025 · 2efc815 · 2efc815
1 parent d4b3756
commit 2efc815
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git a/xla/service/spmd/shardy/round_trip_common/BUILD b/xla/service/spmd/shardy/round_trip_common/BUILD
@@ -22,6 +22,7 @@ cc_library(
         "//xla/service/spmd/shardy:constants",
         "//xla/service/spmd/shardy:utils",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",

diff --git a/xla/service/spmd/shardy/round_trip_common/import_sdy_custom_calls.cc b/xla/service/spmd/shardy/round_trip_common/import_sdy_custom_calls.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/log/check.h"
+#include "absl/strings/string_view.h"
 #include "llvm/ADT/StringRef.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -60,9 +61,10 @@ mlir::LogicalResult rewriteShardingCustomCall(
 
   std::vector<int64_t> unspecDims;
   if (std::optional<mlir::Attribute> backendConfig = op.getBackendConfig()) {
+    StringRef configStr =
+        mlir::dyn_cast<mlir::StringAttr>(*backendConfig).getValue();
     CHECK_OK(xla::sharding_op_util::ParseAttributes(
-        mlir::dyn_cast<mlir::StringAttr>(*backendConfig).getValue().str(),
-        &unspecDims));
+        absl::string_view(configStr.data(), configStr.size()), &unspecDims));
   }
 
   if (op->getNumResults() != 1) {

diff --git a/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc b/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc
@@ -77,8 +77,7 @@ class ManualComputationPattern : public OpConversionPattern<CallOp> {
   mlir::LogicalResult matchAndRewrite(
       CallOp callOp, OpAdaptor adaptor,
       mlir::ConversionPatternRewriter& rewriter) const override {
-    if (!absl::StrContains(callOp.getCallee().str(),
-                           kManualComputationBodyFuncName.str())) {
+    if (!callOp.getCallee().contains(kManualComputationBodyFuncName)) {
       return mlir::failure();
     }
 
@@ -159,8 +158,7 @@ class SdyRoundTripShardMapImportPass
     MLIRContext& context = getContext();
     mlir::ConversionTarget target(context);
     target.addDynamicallyLegalOp<CallOp>([](CallOp op) {
-      return !absl::StrContains(op.getCallee().str(),
-                                kManualComputationBodyFuncName.str());
+      return !op.getCallee().contains(kManualComputationBodyFuncName);
     });
     target.addLegalOp<sdy::ManualComputationOp, sdy::ReturnOp, CustomCallOp>();
     mlir::RewritePatternSet patterns(&context);

diff --git a/xla/service/spmd/shardy/utils.h b/xla/service/spmd/shardy/utils.h
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "absl/strings/escaping.h"
+#include "absl/strings/string_view.h"
 #include "mlir/AsmParser/AsmParser.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Attributes.h"
@@ -80,14 +81,16 @@ template <typename AttrTy>
 AttrTy parseStringAttr(mlir::DictionaryAttr dictAttr,
                        llvm::StringRef attrName) {
   if (mlir::Attribute stringAttr = dictAttr.get(attrName)) {
-    std::string value;
+    std::string unescapedValue;
     std::string error;
+    llvm::StringRef escapedValue =
+        mlir::cast<mlir::StringAttr>(stringAttr).getValue();
     CHECK(absl::CUnescape(
-        mlir::cast<mlir::StringAttr>(stringAttr).getValue().str(), &value,
-        &error))
+        absl::string_view(escapedValue.data(), escapedValue.size()),
+        &unescapedValue, &error))
         << error;
     return mlir::cast<AttrTy>(
-        mlir::parseAttribute(value, stringAttr.getContext()));
+        mlir::parseAttribute(unescapedValue, stringAttr.getContext()));
   }
   return nullptr;
 }