[IFRT] Modify IFRT <-> VIFRT legalization to support escaped SymbolRe…

…fAttr.

PiperOrigin-RevId: 700035710

xla/python/ifrt/ir/tests/ifrt_remove_attrs_from_other_dialects.mlir

@@ -0,0 +1,30 @@
+// RUN: ifrt-opt %s -ifrt-remove-attrs-from-other-dialects -split-input-file | FileCheck %s
+
+!array = !ifrt.array<tensor<2x2xi32>, #ifrt.sharding_param<1x1 to [0] on 2>,
+                     [0,1]>
+// CHECK-LABEL: @non_ifrt_or_builtin_attrives_removed
+// CHECK-NOT: mhlo
+// CHECK: jax.buffer_donor
+// CHECK: xla_tpu_user_reserved_hbm_bytes
+module @non_ifrt_or_builtin_attrives_removed attributes {
+    mhlo.num_partitions = 4 : i32,
+    mhlo.num_replicas = 1 : i32,
+    mpmd.fragments.global_view} {
+  func.func @main(%arg0: !array {mhlo.sharding = "{replicated}"} loc("w1"))
+      -> (!array {mhlo.other_info = ""}) attributes {ifrt.function} {
+    %0, %ctrl = ifrt.Call @add_one::@main(%arg0) on devices [0,1] {
+      ifrt.mesh_name = "mesh1",
+      mhlo.smuggle_attr} : (!array) -> !array
+    return %0 : !array loc("sum_w1")
+  }
+
+  module @add_one attributes {sym_visibility = "private"} {
+    func.func @main(%arg0: tensor<2x2xi32> {jax.buffer_donor = true})
+        -> tensor<2x2xi32> attributes {
+          xla_tpu_user_reserved_hbm_bytes = 0 : i64} {
+      %0 = stablehlo.constant dense<1> : tensor<2x2xi32>
+      %1 = stablehlo.add %arg0, %0 : tensor<2x2xi32>
+      return %1 : tensor<2x2xi32>
+    }
+  }
+}

xla/python/ifrt/ir/tests/vifrt/ifrt_legalize_to_vifrt.mlir

@@ -355,6 +355,19 @@ func.func @op_call(
   %1, %ctrl_1 = ifrt.Call @add_one::@main(%0) after %ctrl_0 on devices [0,1]
       : (!array_op_call) -> !array_op_call
 
+  // Verifies that escaped symbol attr is correctly handled.
+  // CHECK: %[[OUT2:.+]]:2 = "vifrt.CallV1"(%[[ARG0]])
+  // CHECK-SAME: <{
+  // CHECK-DAG: callee = "@escaped-module::@main"
+  // CHECK-DAG: devices = #vifrt<devices_v1[0, 1]>
+  // CHECK-DAG: donated_input_indices = array<i32>
+  // CHECK-DAG: io_aliases = []
+  // CHECK-DAG: operandSegmentSizes = array<i32: 1, 0>
+  // CHECK-SAME: }>
+  // CHECK-SAME: (!vifrt.array_v1<tensor<2x2xi32>, #vifrt.sharding_param_v1<2x1 to [0] on 2>, [0, 1], memory_kind = "vifrt.default">) -> (!vifrt.array_v1<tensor<2x2xi32>, #vifrt.sharding_param_v1<2x1 to [0] on 2>, [0, 1], memory_kind = "vifrt.default">, !vifrt.control_v1)
+  %2, %ctrl_2 = ifrt.Call @"escaped-module"::@main(%arg0) on devices [0,1]
+      : (!array_op_call) -> !array_op_call
+
   // Verifies that the donated input indices attribute is converted.
 
   // CHECK: "vifrt.CallV1"(%[[ARG0]])
@@ -366,7 +379,7 @@ func.func @op_call(
   // CHECK-DAG: operandSegmentSizes = array<i32: 1, 0>
   // CHECK-SAME: }>
   // CHECK-SAME: (!vifrt.array_v1<tensor<2x2xi32>, #vifrt.sharding_param_v1<2x1 to [0] on 2>, [0, 1], memory_kind = "vifrt.default">) -> (!vifrt.array_v1<tensor<2x2xi32>, #vifrt.sharding_param_v1<2x1 to [0] on 2>, [0, 1], memory_kind = "vifrt.default">, !vifrt.control_v1)
-  %2, %ctrl_2 = ifrt.Call @add_one::@main(%arg0) on devices [0,1]
+  %3, %ctrl_3 = ifrt.Call @add_one::@main(%arg0) on devices [0,1]
       {donated_input_indices=array<i32: 0>} : (!array_op_call) -> !array_op_call
 
   // Verifies that the io_aliases attribute is converted.
@@ -380,7 +393,7 @@ func.func @op_call(
   // CHECK-DAG: operandSegmentSizes = array<i32: 1, 0>
   // CHECK-SAME: }>
   // CHECK-SAME: (!vifrt.array_v1<tensor<2x2xi32>, #vifrt.sharding_param_v1<2x1 to [0] on 2>, [0, 1], memory_kind = "vifrt.default">) -> (!vifrt.array_v1<tensor<2x2xi32>, #vifrt.sharding_param_v1<2x1 to [0] on 2>, [0, 1], memory_kind = "vifrt.default">, !vifrt.control_v1)
-  %3, %ctrl_3 = ifrt.Call @add_two::@main(%arg1) on devices [0,1]
+  %4, %ctrl_4 = ifrt.Call @add_two::@main(%arg1) on devices [0,1]
       {io_aliases=[array<i32: 0, 0>]} : (!array_op_call) -> !array_op_call
 
   return %1 : !array_op_call
@@ -395,6 +408,15 @@ module @add_one attributes {sym_visibility = "private"} {
   }
 }
 
+// CHECK-NOT @"escaped-module"
+module @"escaped-module" attributes {sym_visibility = "private"} {
+  func.func private @main(%arg0: tensor<2x2xi32>) -> tensor<2x2xi32> {
+    %0 = stablehlo.constant dense<2> : tensor<2x2xi32>
+    %1 = stablehlo.add %arg0, %0 : tensor<2x2xi32>
+    return %1 : tensor<2x2xi32>
+  }
+}
+
 // CHECK-NOT @add_two
 module @add_two attributes {sym_visibility = "private"} {
   func.func private @main(%arg0: tensor<2x2xi32>) -> tensor<2x2xi32> {

xla/python/ifrt/ir/transforms/BUILD

@@ -41,6 +41,7 @@ cc_library(
         "ifrt_outline_atom_program_to_module_pass.cc",
         "ifrt_populate_atom_program_metadata_pass.cc",
         "ifrt_precompile_atom_program_preprocessing_pass.cc",
+        "ifrt_remove_attrs_from_other_dialects_pass.cc",
         "ifrt_remove_ifrt_attrs_pass.cc",
         "ifrt_reshard_to_copy_arrays_pass.cc",
         "ifrt_verify_bound_external_loaded_executable_pass.cc",

xla/python/ifrt/ir/transforms/ifrt_legalize_to_vifrt_pass.cc

@@ -343,14 +343,17 @@ class IfrtToVifrtOpConverter : public mlir::OpConversionPattern<IfrtOpTy> {
     llvm::DenseSet<mlir::StringAttr> already_converted_attrs;
     if constexpr (std::is_same<IfrtOpTy, CallOp>::value) {
       auto call_op = static_cast<CallOp>(ifrt_op);
-      // Convert the callee from SymbolRefAttr to SymbolNameAttr so that DCE
+      // Convert the callee from SymbolRefAttr to StringAttr so that DCE
       // can remove the atom programs, which have independently legalized to
-      // VHLO.
-      std::string symbol_ref_str;
-      {
-        llvm::raw_string_ostream os(symbol_ref_str);
-        call_op.getCalleeAttr().print(os);
-      }
+      // VHLO. Manually to the conversion by merging RootReference and
+      // NestedReferences to avoid string escaping.
+      std::string symbol_ref_str = absl::StrCat(
+          "@", call_op.getCalleeAttr().getRootReference().getValue().str(),
+          absl::StrJoin(
+              call_op.getCalleeAttr().getNestedReferences(), "",
+              [](std::string* out, const mlir::FlatSymbolRefAttr& symbol_ref) {
+                absl::StrAppend(out, "::@", symbol_ref.getValue().str());
+              }));
       vifrt_attrs.push_back(
           {call_op.getCalleeAttrName(),
            mlir::StringAttr::get(call_op.getContext(), symbol_ref_str)});

xla/python/ifrt/ir/transforms/ifrt_remove_attrs_from_other_dialects_pass.cc

@@ -0,0 +1,160 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "xla/python/ifrt/ir/constants.h"
+#include "xla/python/ifrt/ir/ifrt_dialect.h"
+#include "xla/python/ifrt/ir/transforms/passes.h"
+
+namespace xla {
+namespace ifrt {
+
+namespace {
+
+#define GEN_PASS_DEF_IFRTREMOVEATTRSFROMOTHERDIALECTSPASS
+#include "xla/python/ifrt/ir/transforms/passes.h.inc"
+
+class IfrtRemoveAttrsFromOtherDialectsPass
+    : public impl::IfrtRemoveAttrsFromOtherDialectsPassBase<
+          IfrtRemoveAttrsFromOtherDialectsPass> {
+ public:
+  void runOnOperation() override;
+};
+
+// Returns true if the given `NamedAttribute` is from the IFRT or builtin
+// dialect.
+bool isBuiltinOrIfrtAttr(mlir::NamedAttribute attr) {
+  if (!attr.getNameDialect()) {
+    return true;
+  }
+  auto dialect_namespace = attr.getNameDialect()->getNamespace();
+  return dialect_namespace == mlir::BuiltinDialect::getDialectNamespace() ||
+         dialect_namespace == IfrtDialect::getDialectNamespace();
+}
+
+bool isBuiltinOrIfrtAttr(mlir::Attribute attr) {
+  auto dialect_namespace = attr.getDialect().getNamespace();
+  return dialect_namespace == mlir::BuiltinDialect::getDialectNamespace() ||
+         dialect_namespace == IfrtDialect::getDialectNamespace();
+}
+
+// Returns true if the given `Operation` is an IFRT op, or if it is a FuncOp or
+// ReturnOp of an IFRT function.
+bool isIfrtOpOrFunc(mlir::Operation* op) {
+  if (op->getDialect()->getNamespace() == IfrtDialect::getDialectNamespace()) {
+    return true;
+  }
+  if (auto func_op = llvm::dyn_cast_or_null<mlir::func::FuncOp>(op)) {
+    return op->hasAttr(kIfrtFunctionAttrName);
+  } else if (auto return_op =
+                 llvm::dyn_cast_or_null<mlir::func::ReturnOp>(op)) {
+    return return_op->getParentOp()->hasAttr(kIfrtFunctionAttrName);
+  }
+  return false;
+}
+
+// Recursively removes attributes from the given `Attribute` that are not from
+// the IFRT or builtin dialect.
+mlir::FailureOr<mlir::Attribute> removeAttrsFromOtherDialects(
+    mlir::Attribute attr) {
+  // Remove invalid attributes from container attributes.
+  if (auto array_attr = llvm::dyn_cast<mlir::ArrayAttr>(attr)) {
+    llvm::SmallVector<mlir::Attribute> elements;
+    for (auto element : array_attr.getValue()) {
+      if (auto converted_element_or = removeAttrsFromOtherDialects(element);
+          mlir::succeeded(converted_element_or)) {
+        elements.push_back(*converted_element_or);
+      }
+    }
+    auto res = mlir::ArrayAttr::get(attr.getContext(), elements);
+    return res;
+  } else if (auto dict_attr = llvm::dyn_cast<mlir::DictionaryAttr>(attr)) {
+    llvm::SmallVector<mlir::NamedAttribute> kept_attrs;
+    for (auto named_attr : dict_attr.getValue()) {
+      if (isBuiltinOrIfrtAttr(named_attr)) {
+        if (auto new_attr_or =
+                removeAttrsFromOtherDialects(named_attr.getValue());
+            mlir::succeeded(new_attr_or)) {
+          kept_attrs.push_back(
+              mlir::NamedAttribute(named_attr.getName(), *new_attr_or));
+        }
+      }
+    }
+    auto res = mlir::DictionaryAttr::get(attr.getContext(), kept_attrs);
+    return res;
+  }
+  if (isBuiltinOrIfrtAttr(attr)) {
+    return attr;
+  }
+  return mlir::failure();
+}
+
+mlir::LogicalResult removeAttrsFromOtherDialects(mlir::Operation* op) {
+  auto attr_dict_or = removeAttrsFromOtherDialects(op->getAttrDictionary());
+  if (mlir::succeeded(attr_dict_or)) {
+    if (auto attr_dict =
+            llvm::dyn_cast_or_null<mlir::DictionaryAttr>(*attr_dict_or)) {
+      op->setAttrs(attr_dict);
+    } else {
+      op->emitOpError() << "Failed to remove attrs from other dialects. Remove "
+                           "returned a non-dictionary attribute";
+      return mlir::failure();
+    }
+  } else {
+    op->setAttrs(mlir::DictionaryAttr::get(op->getContext(), {}));
+  }
+  return mlir::success();
+}
+
+void IfrtRemoveAttrsFromOtherDialectsPass::runOnOperation() {
+  mlir::ModuleOp module_op = getOperation();
+  if (mlir::failed(removeAttrsFromOtherDialects(module_op))) {
+    signalPassFailure();
+    return;
+  }
+  auto result = module_op.walk([&](mlir::Operation* op) {
+    if (isIfrtOpOrFunc(op) && mlir::failed(removeAttrsFromOtherDialects(op))) {
+      return mlir::WalkResult::interrupt();
+    }
+    return mlir::WalkResult::advance();
+  });
+  if (result.wasInterrupted()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateIfrtRemoveAttrsFromOtherDialectsPass() {
+  return std::make_unique<IfrtRemoveAttrsFromOtherDialectsPass>();
+}
+
+}  // namespace ifrt
+}  // namespace xla

xla/python/ifrt/ir/transforms/passes.cc

@@ -67,6 +67,7 @@ void CreateIfrtToVersionedPipeline(mlir::OpPassManager& pm,
                                    std::string ifrt_target_version,
                                    std::string vhlo_target_version,
                                    IfrtIrProgramProto& ifrt_ir_program) {
+  pm.addPass(CreateIfrtRemoveAttrsFromOtherDialectsPass());
   pm.addPass(CreateIfrtAtomProgramsToVhloPass(
       ifrt_ir_program.mutable_atom_programs(), std::move(vhlo_target_version)));
   pm.addPass(createIfrtLegalizeToVifrtPass());

xla/python/ifrt/ir/transforms/passes.h

@@ -72,6 +72,9 @@ CreateIfrtLowerAtomProgramMetadataToXlaPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 CreateIfrtRemoveIfrtAttrsPass();
 
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateIfrtRemoveAttrsFromOtherDialectsPass();
+
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 CreateIfrtLowerMpmdReshardToCallPass();
 

xla/python/ifrt/ir/transforms/passes.td

@@ -374,6 +374,16 @@ module {
   let constructor = "CreateIfrtRemoveIfrtAttrsPass()";
 }
 
+def IfrtRemoveAttrsFromOtherDialectsPass :
+    Pass<"ifrt-remove-attrs-from-other-dialects", "mlir::ModuleOp"> {
+  let summary = "Remove attrs from other dialects in IFRT functions";
+  let description = [{
+Remove attributes that are not from IFRT or Builtin dialect from IFRT functions.
+Does not modify atom programs.
+  }];
+  let constructor = "CreateIfrtRemoveAttrsFromOtherDialectsPass()";
+}
+
 def IfrtLowerMpmdReshardToCallPass :
     Pass<"ifrt-lower-mpmd-reshard-to-call", "mlir::ModuleOp"> {
   let summary = "Lowers MPMD `ifrt.Reshard` to `ifrt.Call`";

xla/python/ifrt/ir/transforms/vifrt_legalize_to_ifrt_pass.cc

@@ -311,8 +311,10 @@ mlir::FailureOr<mlir::SymbolRefAttr> getCalleeSymbolRef(CallOpV1 call_op) {
   if (!callee_symbol_ref_str_attr) {
     return mlir::failure();
   }
-  std::vector<std::string> symbol_strs =
-      absl::StrSplit(callee_symbol_ref_str_attr.str(), absl::ByString("::@"));
+  // It is important to call `getValue()` on the `StringAttr` to get the
+  // unescaped string instead of the escaped string.
+  std::vector<std::string> symbol_strs = absl::StrSplit(
+      callee_symbol_ref_str_attr.getValue().str(), absl::ByString("::@"));
   if (symbol_strs.empty()) {
     return mlir::failure();
   }