[opt] Unconditionally remove sends' data if their predicates are false

Previously, when we are unable to remove an unconditionally-disabled send, we did nothing; this left the sends marked as using their data, which could prevent other optimizations from kicking in. Now, if we can't remove a send, we replace its data input with a zero-valued literal.

PiperOrigin-RevId: 809231139

Author: ericastor

xls/passes/useless_io_removal_pass.cc

@@ -88,10 +88,23 @@ absl::StatusOr<bool> UselessIORemovalPass::RunInternal(
         }
         XLS_ASSIGN_OR_RETURN(ChannelRef channel_ref, send->GetChannelRef());
         Node* predicate = send->predicate().value();
-        if (query_engine.IsAllZeros(predicate) &&
-            channel_maps.to_send.at(channel_ref).size() >= 2) {
-          channel_maps.to_send.at(channel_ref).erase(send);
-          replacement = send->token();
+        if (query_engine.IsAllZeros(predicate)) {
+          // We can remove the send if this is not the last send left on the
+          // channel.
+          if (channel_maps.to_send.at(channel_ref).size() > 1) {
+            channel_maps.to_send.at(channel_ref).erase(send);
+            replacement = send->token();
+          } else if (!send->data()->Is<Literal>() ||
+                     !send->data()->As<Literal>()->value().IsAllZeros()) {
+            // If we aren't removing the send, at least replace its data
+            // input with zero, as it will never be used.
+            XLS_ASSIGN_OR_RETURN(
+                Literal * zero,
+                proc->MakeNode<Literal>(send->loc(),
+                                        ZeroOfType(send->data()->GetType())));
+            send->ReplaceOperand(send->data(), zero);
+            changed = true;
+          }
         } else if (query_engine.IsAllOnes(predicate)) {
           XLS_ASSIGN_OR_RETURN(
               replacement,
@@ -108,6 +121,8 @@ absl::StatusOr<bool> UselessIORemovalPass::RunInternal(
         Node* predicate = receive->predicate().value();
         if (query_engine.IsAllZeros(predicate) &&
             channel_maps.to_receive.at(channel_ref).size() >= 2) {
+          // We can remove the receive if this is not the last receive left on
+          // the channel.
           XLS_ASSIGN_OR_RETURN(Channel * channel, GetChannelUsedByNode(node));
           channel_maps.to_receive.at(channel_ref).erase(receive);
           XLS_ASSIGN_OR_RETURN(Literal * zero,

xls/passes/useless_io_removal_pass_test.cc

@@ -32,7 +32,9 @@
 #include "xls/ir/function_builder.h"
 #include "xls/ir/ir_matcher.h"
 #include "xls/ir/ir_test_base.h"
+#include "xls/ir/nodes.h"
 #include "xls/ir/package.h"
+#include "xls/ir/proc.h"
 #include "xls/ir/source_location.h"
 #include "xls/ir/value.h"
 #include "xls/passes/dce_pass.h"
@@ -77,27 +79,18 @@ TEST_F(UselessIORemovalPassTest, DontRemoveOnlySend) {
       p->CreateStreamingChannel("test_channel", ChannelOps::kSendOnly,
                                 p->GetBitsType(32)));
   ProcBuilder pb(TestName(), p.get());
+  BValue token = pb.StateElement("tkn", Value::Token());
   pb.StateElement("state", Value(UBits(0, 0)));
-  pb.SendIf(channel, pb.Literal(Value::Token()), pb.Literal(UBits(0, 1)),
-            pb.Literal(UBits(1, 32)));
-  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({pb.Literal(UBits(0, 0))}));
-  int64_t original_node_count = proc->node_count();
-  EXPECT_THAT(Run(p.get()), IsOkAndHolds(false));
-  EXPECT_EQ(proc->node_count(), original_node_count);
-}
-
-TEST_F(UselessIORemovalPassTest, DontRemoveOnlySendNewStyle) {
-  auto p = CreatePackage();
-  TokenlessProcBuilder pb(NewStyleProc(), TestName(), "tkn", p.get());
-  XLS_ASSERT_OK_AND_ASSIGN(
-      SendChannelInterface * channel,
-      pb.AddOutputChannel("test_channel", p->GetBitsType(32)));
-  pb.StateElement("state", Value(UBits(0, 0)));
-  pb.SendIf(channel, pb.Literal(UBits(0, 1)), pb.Literal(UBits(1, 32)));
-  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({pb.Literal(UBits(0, 0))}));
+  token = pb.SendIf(channel, token, pb.Literal(UBits(0, 1)),
+                    pb.Literal(UBits(1, 32)), SourceInfo(), "my_send");
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc,
+                           pb.Build({token, pb.Literal(UBits(0, 0))}));
   int64_t original_node_count = proc->node_count();
-  EXPECT_THAT(Run(p.get()), IsOkAndHolds(false));
+  EXPECT_THAT(Run(p.get()), IsOkAndHolds(true));
   EXPECT_EQ(proc->node_count(), original_node_count);
+  XLS_ASSERT_OK_AND_ASSIGN(Node * send_node, proc->GetNode("my_send"));
+  ASSERT_TRUE(send_node->Is<Send>());
+  EXPECT_THAT(send_node->As<Send>()->data(), m::Literal(0));
 }
 
 TEST_F(UselessIORemovalPassTest, RemoveSendIfLiteralFalse) {
@@ -243,6 +236,48 @@ TEST_F(UselessIORemovalPassTest, RemoveReceivePredIfLiteralTrue) {
       ElementsAre(m::Next(proc->GetStateRead(1), m::TupleIndex(tuple, 1))));
 }
 
+TEST_F(UselessIORemovalPassTest, DontRemoveLastSendIfOnSendOnlyChannel) {
+  auto p = CreatePackage();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      StreamingChannel * channel,
+      p->CreateStreamingChannel("test_channel", ChannelOps::kSendOnly,
+                                p->GetBitsType(32)));
+  ProcBuilder pb(TestName(), p.get());
+  BValue token = pb.StateElement("tkn", Value::Token());
+  pb.StateElement("state", Value(UBits(0, 0)));
+  token = pb.SendIf(channel, token, pb.Literal(UBits(0, 1)),
+                    pb.Literal(UBits(1, 32)), SourceInfo(), "my_send");
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc,
+                           pb.Build({token, pb.Literal(UBits(0, 0))}));
+  ASSERT_EQ(p->channels().size(), 1);
+  int64_t original_node_count = proc->node_count();
+  EXPECT_THAT(Run(p.get()), IsOkAndHolds(true));
+  EXPECT_EQ(proc->node_count(), original_node_count);
+  XLS_ASSERT_OK_AND_ASSIGN(Node * send_node, proc->GetNode("my_send"));
+  ASSERT_TRUE(send_node->Is<Send>());
+  EXPECT_THAT(send_node->As<Send>()->data(), m::Literal(0));
+}
+
+TEST_F(UselessIORemovalPassTest, DontRemoveLastReceiveIfOnReceiveOnlyChannel) {
+  auto p = CreatePackage();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      StreamingChannel * channel,
+      p->CreateStreamingChannel("test_channel", ChannelOps::kReceiveOnly,
+                                p->GetBitsType(32)));
+  ProcBuilder pb(TestName(), p.get());
+  BValue token = pb.StateElement("tkn", Value::Token());
+  pb.StateElement("state", Value(UBits(0, 32)));
+  BValue token_and_result =
+      pb.ReceiveIf(channel, token, pb.Literal(UBits(0, 1)));
+  token = pb.TupleIndex(token_and_result, 0);
+  BValue result = pb.TupleIndex(token_and_result, 1);
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({token, result}));
+  ASSERT_EQ(p->channels().size(), 1);
+  int64_t original_node_count = proc->node_count();
+  EXPECT_THAT(Run(p.get()), IsOkAndHolds(false));
+  EXPECT_EQ(proc->node_count(), original_node_count);
+}
+
 void IrFuzzUselessIORemoval(FuzzPackageWithArgs fuzz_package_with_args) {
   UselessIORemovalPass pass;
   OptimizationPassChangesOutputs(std::move(fuzz_package_with_args), pass);