[IFRT] Add pass to remove attributes that are not from IFRT or Builti…

…n dialects.

This pass runs as part of the versioning pipeline, and its purpose is to prevent smuggling of attributes that are not from the Builtin or IFRT dialect. Such attributes might be present as result of JAX lowering (e.g., mhlo.sharding), but IFRT IR should not use them nor depend on them. By adding this pass, we avoid failing during IFRT -> VIFRT legalization pass, which fails if such attributes exist.

PiperOrigin-RevId: 700526071

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_attributes_removed
+// CHECK-NOT: mhlo
+// CHECK: jax.buffer_donor
+// CHECK: xla_tpu_user_reserved_hbm_bytes
+module @non_ifrt_or_builtin_attributes_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/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_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`";