Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[IFRT] Add pass to remove attributes that are not from IFRT or Builtin dialects. #19861

Merged
merged 1 commit into from
Nov 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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>
}
}
}
1 change: 1 addition & 0 deletions xla/python/ifrt/ir/transforms/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions xla/python/ifrt/ir/transforms/passes.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
3 changes: 3 additions & 0 deletions xla/python/ifrt/ir/transforms/passes.h
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down
10 changes: 10 additions & 0 deletions xla/python/ifrt/ir/transforms/passes.td
Original file line number Diff line number Diff line change
Expand Up @@ -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`";
Expand Down