remove block args that were added (hacky)

Author: srcarroll

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

@@ -226,17 +226,21 @@ static void generateFusedElementwiseOpRegion(
   // Since some ops, like `linalg.map`, do not have block arguments for init
   // operands then we first "generalize" the block by adding arguments for init
   // operands when they aren't present. We detect this case by checking if
-  // `getOpOperandsMatchingBBargs() == getDpsInputOperands()
+  // `getOpOperandsMatchingBBargs() == getDpsInputOperands()`.
+  // TODO: This is hacky and should not be merged. Keeping for now for testing
+  // purposes in the meantime, but need a better way
   Block &producerBlock = producer->getRegion(0).front();
-  if (producer.getOpOperandsMatchingBBargs() ==
-      producer.getDpsInputOperands()) {
+  bool addOutputArgsProducer =
+      producer.getOpOperandsMatchingBBargs() == producer.getDpsInputOperands();
+  if (addOutputArgsProducer) {
     for (auto init : producer.getDpsInits())
       producerBlock.addArgument(getElementTypeOrSelf(init.getType()),
                                 producer.getLoc());
   }
   Block &consumerBlock = consumer->getRegion(0).front();
-  if (consumer.getOpOperandsMatchingBBargs() ==
-      consumer.getDpsInputOperands()) {
+  bool addOutputArgsConsumer =
+      consumer.getOpOperandsMatchingBBargs() == consumer.getDpsInputOperands();
+  if (addOutputArgsConsumer) {
     for (auto init : consumer.getDpsInits())
       consumerBlock.addArgument(getElementTypeOrSelf(init.getType()),
                                 consumer.getLoc());
@@ -350,6 +354,14 @@ static void generateFusedElementwiseOpRegion(
   // Sanity checks.
   assert(fusedBlock->getNumArguments() == fusedOp->getNumOperands() &&
          "Ill-formed GenericOp region");
+  // Erase added args in case that the ops are still live after fusion.
+  // TODO: Remove along with hacky code above.
+  if (addOutputArgsProducer)
+    producerBlock.eraseArguments(producer.getNumDpsInputs(),
+                                 producer.getNumDpsInits());
+  if (addOutputArgsConsumer)
+    consumerBlock.eraseArguments(consumer.getNumDpsInputs(),
+                                 consumer.getNumDpsInits());
 }
 
 FailureOr<mlir::linalg::ElementwiseOpFusionResult>