[WIP] Support Alias V2

[eedalong]

No description provided.

[Yancey1989]

what's the difference with ArgsMutationOp rewrite pass? I prefer using or extend args-mutation op instead of a new AliasOp.

BladeDISC/tao_compiler/mlir/disc/transforms/disc_lhlo_rewriter.cc

Lines 69 to 82 in ddcc491

	struct LhloArgsMutationOpRewriter
	: public OpRewritePattern<lmhlo_disc::ArgsMutationOp> {
	explicit LhloArgsMutationOpRewriter(MLIRContext* context)
	: OpRewritePattern(context) {}
	LogicalResult matchAndRewrite(lmhlo_disc::ArgsMutationOp lhloOp,
	PatternRewriter& rewriter) const override {
	auto op = lhloOp.getOperation();
	auto operands = op->getOperands();
	Value value = backtraceOperand<memref::ReinterpretCastOp>(operands[0]);
	value.replaceAllUsesWith(operands[1]);
	rewriter.eraseOp(op);
	return success();
	}
	};

[eedalong]

This is a WIP version, in the formal version, we need to add inc_ref && dealloc pair for safe memory reuse. But here since we are reusing module input, we can be sure these input tensors are alive during DISC execution.

[eedalong]

Yeah, extending ArgsMutationOp is definitely Ok. But when we try to reuse buffer which is not one of the op input, we might want a better name to represent this kind of situation.