Add support for executable duplication in encoding specialization pass.

[hanhanW]

Duplicates stream.executables based on the affinity analysis of stream.async.dispatch ops. Some executables can be launched by different devices. It can produce wrong codegen artifacts when bindings types are encoded (i.e., the tensor type has an encoding attribute). Because they can result in different layouts, especially when multi-device is involved. E.g., say that device_a and device_b interpret a tensor type with encodings in different layouts, and there is an executable that can be launch with resources from either device_a or device_b. It is confusing what the input layouts for the executable because there are two possibilities. In this case, we have to duplicate the executable with updated encoding, and modify the dispatch to launch proper executable based on device analysis.

[hanhanW]

The PR depends on #19502. It is ready for review, which compares the diff to the parent mirror branch.

[hanhanW]

I spent some time on the bazel failure, but I can't reproduce it on cmake build even with ASAN enabled. I'll try to set up the bazel one my env, the PR is ready to review. Please take a look, thanks!

[hanhanW]

I spent some time on the bazel failure, but I can't reproduce it on cmake build even with ASAN enabled. I'll try to set up the bazel one my env, the PR is ready to review. Please take a look, thanks!

I think I'd be able to reproduce it if I build it in release mode. I put a statement in an assertion because I wanted to check if the return value is true or not. However, the statement is dropped in release build..

[benvanik]

We don't want to duplicate based on affinity - in your test we can see why: we don't want to double compile executables that are going to have the same exact compilation flow as we're just doubling our compile time (or on an 8 way config x8!).

IREE::Stream::AffinityAttr has a areCompatible / canExecuteTogether which indicates when two operations can execute in the same place (which is not what you want - that's for execution scheduling). Here you need something like areTranslationCompatible which indicates when two affinities share the same execution configuration and only duplicate when that's not the case (you want the smallest set of unique executable targets). We could probably rename areCompatible to areExecutionCompatible to differentiate (at some point we'll add an areAllocationCompatible for unified memory systems and stuff where execution and translation may differ but buffers can be shared).

Note to make the compatibility checks scale better you can gather all (like you do), use llvm::EquivalenceClasses + areTranslationCompatible to build the sets of compatible devices, then run down your executables and duplicate for each non-intersecting set (this avoids doing areTranslationCompatible n^2).

You'll want to have a test like you have with two devices that share targets ensuring they don't get duplicated and two devices with different targets ensuring they do.

[benvanik]

style: prefer not using numbers - they always end up wrong when code is moved around or when code is inserted/removed. You can call out preconditions/postconditions if blocks of code do have to be performed in order (e.g. "build cache of valid ops after removing dead ops" etc)

if the algorithm has a flow then that can be documented in the function comment where someone landing on the function wants to know how it works. By keeping them inline a reader has to read the whole function (like I just did, wondering where 2 was for a page :)

[hanhanW]

iree/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td

Lines 108 to 116 in 36c2353

	// TODO(benvanik): replace with more fine-grained compatibility checks.
	// "Compatible" can mean a lot of things: are they cache-coherent, are they
	// a shared address space, are they able to perform collective operations,
	// etc. May be able to represent it with a bitfield or a dedicated
	// compatibility struct result.
	// Returns true if |desiredAffinity| (if any) is compatible with
	// |requiredAffinity|.
	static bool areCompatible(AffinityAttr desiredAffinity,
	AffinityAttr requiredAffinity);

I found the above comment in areCompatible documentation, and I have a question. It seems like it'd return true if one of desiredAffinity is compatible with requiredAffinity. Does it mean that [device_a, device_b] is compatible with [device_a]? Does it mean that the compatibility could be unidirectional? I.e., X is compatible with Y, but Y is not compatible with X.

In the areTranslationCompatible, should we make it bidirectional?

Note to make the compatibility checks scale better you can gather all (like you do), use llvm::EquivalenceClasses + areTranslationCompatible to build the sets of compatible devices, then run down your executables and duplicate for each non-intersecting set (this avoids doing areTranslationCompatible n^2).

IIUC, the llvm::EquivalenceClasses is a disjoint-set. Say that we have areTranslationCompatible implemented. The algorithm would be:

1. Collect all the affinities.
2. Categorize the affinities into few groups, using llvm::EquivalenceClasses.
3. Each group can map to a unique hash/affinity/whatever. We then use the unique affinity in my current implementation, which can reduce the number of duplications.

Do I understand it correctly?

Meanwhile, I'm going to study the compatibility things in Stream dialect. It's a new thing to me!

[hanhanW]

I just found that the current implementation is simple. It does not consider the case that I listed above. I think I can start with a similar version.

iree/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp

Lines 375 to 384 in 36c2353

	bool AffinityAttr::areCompatible(AffinityAttr desiredAffinity,
	AffinityAttr requiredAffinity) {
	if (desiredAffinity == requiredAffinity)
	return true;
	if ((desiredAffinity && !requiredAffinity) ||
	(requiredAffinity && !desiredAffinity)) {
	return true;
	}
	// We could do a fuzzier match here (interface isCompatible() etc).
	return false;

[benvanik]

the executable targets are what we care about matching, not the affinities - we don't want to cluster by affinities, but by those sets of targets, and then generate based on required sets of targets. CUDA GPU 0 and CUDA GPU 1 (two different affinities) have the same executable targets and we don't want to duplicate, CUDA GPU 0, CUDA GPU 1, and CPU (three different affinities) have CUDA GPU 0 and CUDA GPU 1 sharing the same targets but the CPU is different and we'd want one copy for the CUDA GPUs and one copy for the CPU (because we're talking about targets, not affinities).

[benvanik]

(that may be what you're saying, I just realized I didn't word what I said very well :)