Add a side effect trait for XlaLaunch op

This is a follow up to cl/516624982, which incorrectly used TPUExecute's side effect trait for XlaLaunch ops. PiperOrigin-RevId: 516685949
GCABC123 · Mar 15, 2023 · 594f7e0 · 594f7e0
1 parent b15ecc9
commit 594f7e0
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
@@ -187,6 +187,7 @@ def TF_RecvResource : TF_ResourceBase<"Recv">;
 def TF_TPUExecuteResource : TF_ResourceBase<"TPUExecute">;
 def TF_RandomGeneratorResource : TF_ResourceBase<"RandomGenerator">;
 def TF_XlaHostComputeResource : TF_ResourceBase<"XlaHostCompute">;
+def TF_XlaLaunchResource : TF_ResourceBase<"XlaLaunch">;
 def TF_CollectiveReduceOrderingResource : TF_ResourceBase<"CollectiveReduceOrdering">;
 def TF_NcclAllReduceOrderingResource : TF_ResourceBase<"NcclAllReduceOrdering">;
 
@@ -249,7 +250,7 @@ def TF_SendSideEffect : MemoryEffects<[MemWrite<TF_SendResource>]>;
 def TF_RecvSideEffect : MemoryEffects<[MemWrite<TF_RecvResource>]>;
 def TF_XlaHostComputeSideEffect : MemoryEffects<[MemWrite<TF_XlaHostComputeResource>]>;
 
-def TF_TPUExecuteSideEffect : MemoryEffects<[MemWrite<TF_TPUExecuteResource>]>;
+def TF_XlaLaunchSideEffect : MemoryEffects<[MemWrite<TF_XlaLaunchResource>]>;
 
 def TF_RandomGeneratorSideEffect : MemoryEffects<[MemWrite<TF_RandomGeneratorResource>]>;
 

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc
@@ -2927,7 +2927,7 @@ void XlaLaunchOp::getEffects(
         &effects) {
   effects.reserve(getArgs().size() + 1);
   effects.emplace_back(MemoryEffects::Write::get(),
-                       ResourceEffects::TPUExecute::get());
+                       ResourceEffects::XlaLaunch::get());
 
   for (Value value : getArgs()) {
     if (value.getType()

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h b/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
@@ -104,6 +104,10 @@ struct NcclAllReduceOrdering
   StringRef getName() final { return "NcclAllReduceOrdering"; }
 };
 
+struct XlaLaunch : public ::mlir::SideEffects::Resource::Base<XlaLaunch> {
+  StringRef getName() final { return "XlaLaunch"; }
+};
+
 // Returns true iff resource type with given ID is only self-dependent, i.e.,
 // there are no dependencies to other resource types (including unknown resource
 // type).