remove hack for linalg.map and update tests

Author: srcarroll

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

@@ -224,28 +224,8 @@ static void generateFusedElementwiseOpRegion(
   auto consumer = cast<LinalgOp>(fusedOperand->getOwner());
   // Build the region of the fused op.
 
-  // 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()`.
-  // 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();
-  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();
-  bool addOutputArgsConsumer =
-      consumer.getOpOperandsMatchingBBargs() == consumer.getDpsInputOperands();
-  if (addOutputArgsConsumer) {
-    for (auto init : consumer.getDpsInits())
-      consumerBlock.addArgument(getElementTypeOrSelf(init.getType()),
-                                consumer.getLoc());
-  }
   OpBuilder::InsertionGuard guard(rewriter);
   Block *fusedBlock = rewriter.createBlock(&fusedOp->getRegion(0));
   IRMapping mapper;
@@ -355,14 +335,6 @@ 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>

mlir/test/Dialect/Linalg/fusion-elementwise.mlir

@@ -89,7 +89,7 @@ func.func @map_ops_mixed_types(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> te
   %0 = linalg.map {math.sqrt} ins(%arg0 : tensor<8xf32>) outs(%initf : tensor<8xf32>)
   %1 = linalg.map {math.exp} ins(%arg1 : tensor<8xf32>) outs(%initf : tensor<8xf32>)
   %2 = linalg.map ins(%0, %1 : tensor<8xf32>, tensor<8xf32>) outs (%init : tensor<8xi1>)
-    (%in0 : f32, %in1 : f32) {
+    (%in0 : f32, %in1 : f32, %out : f32) {
       %cmp = arith.cmpf olt, %in0, %in1 : f32
       linalg.yield %cmp : i1
   }
@@ -150,7 +150,7 @@ func.func @elementwise_ops(%in1: tensor<8xf32>, %in2: tensor<8x10xf32>) -> tenso
 func.func @map_multi_ops(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>) -> tensor<8xf32> {
     %fill = tensor.empty() : tensor<8xf32>
     %add_exp = linalg.map ins(%arg0, %arg1: tensor<8xf32>, tensor<8xf32>) outs(%fill: tensor<8xf32>)
-    (%in0 : f32, %in1 : f32) {
+    (%in0 : f32, %in1 : f32, %out : f32) {
       %add = arith.addf %in0, %in1 : f32
       %exp = math.exp %add : f32
       linalg.yield %exp : f32