Implement live variable analysis

[sherhut]

It would be great to have a shared implementation for live variable analysis in core that could be used by various resource allocation schemes in tensorflow and others.

[dfki-mako]

We would like to implement such an analysis. Based on our insights into MLIR we propose the following high-level interface for this new analysis:

namespace mlir {
    class ValueSet;

    class LivenessValueInfo {
    public:
        /// Returns all values that are live at this
        /// point in time. Note that this refers to the
        /// associated operation of the underlying value.
        const ValueSet &live() const;

        /// Returns the last operation the associated value
        /// was seen as live.
        Operation *endOfLife() const;
    };

    class LivenessBlockInfo {
    public:
        /// Returns all values that are live at the beginning
        // of the block.
        const ValueSet &in() const;

        /// Returns all values that are live at the end
        // of the block.
        const ValueSet &out() const;
    };

    class Liveness {
    public:
        /// Creates a new Liveness analysis that computed liveness
        /// information for all associated regions.
        Liveness(Operation *op);

        /// Resolves liveness info (if any) for the given value.
        LivenessValueInfo *getLiveness(Value *value) const;

        /// Resolves liveness info (if any) for the given block.
        LivenessBlockInfo *getLiveness(Block *block) const;
    };
}

This interface should satisfy common use-cases like straight-forward register allocation.

[bondhugula]

Looks good! For 'getLiveness(Value *value)', you could instead have a

LivenessValueInfo *getLivenessAtOp(Operation *op)

to (1) avoid the situation where a Value isn't defined by an op (because it could also be a block argument) and (2) to allow computing this information at ops that have zero results.

In addition, I think you need a 'const LivenessValueInfo/LivenessBlockInfo *' on getLiveness or a void getLiveness(Value *value, LivenessValueInfo &info).

[River707]

@bondhugula What benefit does computing on an Operation with zero results have?

@dfki-mako Without looking at how you intend to implement it(and/or use it) I don't have too many comments at this point, but the main one is that I don't understand the endOfLife method. What do you mean by "last operation"? How do you intend to handle control flow graphs where a value doesn't have a single user that post dominates all others?

There is some previous art of computing SSA liveness in LLVM:
https://github.com/llvm/llvm-project/blob/05da2fe52162c80dfa18aedf70cf73cb11201811/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp#L236

As well as an implementation in IREE(on MLIR):
https://github.com/google/iree/blob/4becc9eb04d9e79ce833b20f6fe7ae9fbcd3c80f/iree/compiler/Dialect/VM/Analysis/ValueLiveness.cpp

I have also written a full SSA liveness analysis for a custom LLVM backend, of which partially inspired the IREE implementation. I'm happy to review any PRs you have, or help in any way I can.

[bondhugula]

@bondhugula What benefit does computing on an Operation with zero results have?

IIUC, the getLiveness method is providing liveness info (live-in's and live-out's) at that program point -- as such, it makes sense for any operation (within the op regions that the Liveness instance was constructed on).

~ Uday

[River707]

Ohh sorry, I missed the comment. At least for register allocation, you want the ranges of the program that a value is live and don't necessarily work at a per-op level.

[dfki-mako]

@bondhugula, @River707 Thank you for your feedback on our first proposal.

@bondhugula getLivenessAtOp(Operation*) could be implemented in favor of the suggested getLiveness(Value*). As you mentioned, getLiveness(Value*)` covers cases where an operation has no result value. This can (depending on the use case) lead to problems when querying liveness information at an operating level. We initially thought (like @River707) that it might not be necessary to query information at this level. However, we follow your argument that it makes sense to query liveness information at a certain point in the program - in other words, at a certain operation. What would be the intended output of this query? We would expect to receive information about the liveness of all live values, including the operands of this operation.

The return of a const pointer (like const LivenessValueInfo/LivenessBlockInfo *) depends on your concept of const correctness. The CallGraph analysis, for example, has a const getter for single nodes CallGraphNode *lookupNode(region *region) const;. However, these nodes seem to be mutable, and so the instance CallGraph seems to be mutable, although we used a const getter. This seems to be intentional, but in constrast to this analysis we would suggest that the Liveness would be immutable. Therefore, it makes sense to return a const LivenessValueInfo*. What do you think?

@River707 endOfLive can be somewhat misleading. This simple functionality would only cover simple cases to get started. To have a rock solid version, we would suggest using an endOfLive(Value*, Operation*) method instead. However, it may be useful to add information about the operand index for fine-grained register allocation.

[bondhugula]

Thanks, everything you suggest sounds good to me. My comment on liveness at a point was due to the comment on your method LivenessValueInfo::live(). Clearly, live ranges are what you may first/usually need for buffer assignment, etc., and so you could just go with an interface for that.