[MLIR][XLA] Remove redundant LHLO CopyOp

[dfki-ehna]

The redundant LHLO CopyOp that copies an allocated buffer to a block argument is removed. The uses of the allocated buffer are replaced with the block argument. The associated Alloc and Dealloc are removed.

[rthadur]

@dfki-ehna can you please include test cases around the new changes ?

[pifon2a]

Thanks!

[rthadur]

@pifon2a can you please merge this change internally.

[pifon2a]

(transmitting @joker-eph): You are dereferencing the user of the allocOp while you already deleted the copyOp, which mean there might be no user left.

Here is a minimal reproducer that will show it:

func @Fusion(%arg0: memref<2x2xf32>) -> memref<2x2xf32> {
%0 = alloc() {temp = true} : memref<2x2xf32>
"xla_lhlo.copy"(%0, %arg0) : (memref<2x2xf32>, memref<2x2xf32>) -> ()
return %0 : memref<2x2xf32>
}

This comes back to my other point about testing: with dedicated testing for this pass, I would have expected this IR to be the first basic test to validate the pass.

[dfki-ehna]

Thanks for the hint. You are right that this code snippet would break our current version. Although this code could exist from a theoretical point of view as an input program, the current implementation of the HLO-to-LHLO-Legalization pass ensures that there will always be a dealloc. Furthermore, it is even worse: currently, the Std.Return-to-LHLO converter expects a final dealloc operation to place a proper CopyOp. The following code snippet is taken from the hlo_legalize_to_lhlo.cc file (starting in line 425):

if (dealloc == nullptr) {
    returnOp.emitOpError()
        << "Missing dealloc for operand " << operand.index();
    return matchFailure();
}

This code snippet also fails in the following tiny test case:

func @TestFunc(%arg0: tensor<2xf32>) -> tensor<2xf32> {
   return %arg0 : tensor<2xf32>
}

[dfki-ehna]

We have addressed all the comments in the new PR.

[joker-eph]

(GitHub does not let me answer inline apparently)

Although this code could exist from a theoretical point of view as an input program, the current implementation of the HLO-to-LHLO-Legalization pass ensures that there will always be a dealloc.

I think this hits a principle of development for MLIR passes: the only invariant you can assume on your input is that the verifier is passing. It isn't allowed to crash on random vali input, we should be able to fuzz your pass safely (you can gracefully signal a pass failure and return though)