[Mlir-commits] [mlir] [mlir][TilingInterface] Allow multiple results in PartialReductionOpInterface (PR #92624)

[Kunwar Grover]

================
@@ -155,12 +155,12 @@ static Value createDestinationPassingStyleInitOperand(
         tensor::getMixedSizes(builder, builder.getLoc(), spmdizedOperand);
     PartialReductionOpInterface partialReductionIface =
         llvm::cast<PartialReductionOpInterface>(op.getOperation());
-    FailureOr<Operation *> reductionNeutralTensorOp =
+    assert(op->getNumResults() == 1 && "Multiple results not supported.");
+    FailureOr<SmallVector<Value>> reductionNeutralTensor =
         partialReductionIface.generateInitialTensorForPartialReduction(
             builder, builder.getLoc(), shape, {});
-    assert(succeeded(reductionNeutralTensorOp));
-    builder.create<scf::YieldOp>(
-        reductionNeutralTensorOp.value()->getResult(0));
+    assert(succeeded(reductionNeutralTensor));
----------------
Groverkss wrote:

I'd rather not change the MeshSharding implementation. I don't know the details and i'd rather just keep it the way it was (it was asserting earlier).

https://github.com/llvm/llvm-project/pull/92624