[LLVMCPU] Fix linking between modules with different target attributes

[bjacob]

Just recording thoughts before holidays...

Around LLVMCPUTarget.cpp, when we load and link different bitcode modules, we are currently just removing some function attributes to make it more or less work, https://github.com/openxla/iree/blob/9cde4e3dc589ab783757b863a747477bc6fafc91/compiler/src/iree/compiler/Dialect/HAL/Target/LLVMCPU/Builtins/UKernel.cpp#L35-L44 , but it's not very rigorous, and it needs to become more rigorous in order to overcome two issues we're having:

• When we remove all function target attributes including target CPU features, and the code inside a function was actually relying on a CPU feature (e.g. was using an intrinsic provided by that CPU feature), we are effectively introducing undefined behavior. It mostly just works, until this week it didn't --- one of the bugs I've dealt with this week was, in int4-quantized Llama2, some avx512vnni LLVM intrinsic was causing a llvm_unreachable to be reached in the LLVM x86 backend.
• What we are currently doing with the attributes-stripping isn't entirely successful at one of its main goals, which is to enable perfect inlining of ukernel code into dispatch functions. It's working on x86 but not on arm64. It's probably a matter of being more rigorous about how we bring attributes in good alignment between caller and callee.

I'm thinking that the solution is, rather than dropping attributes, we just need to add all the attributes from the llvm::TargetMachine. Then, by construction, all functions that are actually legal to call on that target machine will have exactly the same attribute (give or take any preexisting difference in the attributes that they had before). And functions that have additional attributes are not legal to call on this target machine anyway.

As a first approximation to make this concrete, I'm toying with a change like this:

-static void removeTargetAttributes(llvm::Module &module) {
+static void overrideTargetAttributes(llvm::Module &module, const llvm::TargetMachine& targetMachine) {
   for (auto &func : module.functions()) {
-    func.removeFnAttr("target-cpu");
-    func.removeFnAttr("tune-cpu");
-    func.removeFnAttr("target-features");
+    func.addFnAttr("target-cpu", targetMachine.getTargetCPU());
+    func.addFnAttr("target-features", targetMachine.getTargetFeatureString());
   }
 }

I'll get back to this in early 2024... I have been mulling something like this for months as I knew about the issue about ukernels not getting fully inlined on arm64, but now that I know that it's also causing UB manifesting as compiler crashes on Llama2 on x86, this became high priority.

FYI @benvanik @Max191

[benvanik]

yeah definitely too big of a hammer today, but we don't want to replace it with another big hammer. we ideally want to support mixed targets and not override/change anything from certain modules. a single bitcode file should be able to have any number of functions with different cpu/features attributes on them so that we don't explode our deployment/compilation matrix (effectively function multiversioning). I suspect what's there (stripping attrs) and what you propose (changing all attrs) only works because we aren't yet doing this and we'll want to find a solution that doesn't need to be re-solved for that.

the ideal workflow would be to have functions that have a target explicitly specified (either in C before going into bitcode or generated by us) preserve those attrs while functions that are generic take on the attrs of the global target. the thing we'll want to look into is how to identify those functions to preserve/strip from the IR as we walk the functions.

ex:

__attribute__((target("arch=whateverA")))
static void my_avx512_whateverA_thing() { ... }
__attribute__((target("arch=whateverB")))
static void my_avx512_whateverB_thing() { ... }
void my_avx512_base_thing() {
  if (runtime_feature_query_A) my_avx512_whateverA_thing();
  else my_avx512_whateverB_thing();
}

(clang docs say __attribute__((target("default"))) is how to indicate default target, but we'd have to check the bitcode and see if that comes through or if that gets expanded to the default target before serialization)

[bjacob]

We actually already have information on which bitcode functions are generic and need to have their target attributes overwritten: they are in separate bitcode files to begin with (they have to be, as we build them separately) and they also have a different target architecture (wasm). https://github.com/openxla/iree/blob/dc81beb1199d95d800604491c80ade8fbdc2534a/runtime/src/iree/builtins/ukernel/CMakeLists.txt#L132-L148

So a minimal approach to this issue might be, only for these generic ukernel bitcode files, overwrite with the targetMachine's attributes, and leave the architecture-specific bitcode attributes unchanged.

I amn't clear exactly how finely aligned attributes need to be, between caller and callee, to enable inlining.

If we still are missing inlinings with the above, then I still think that what I propose in the issue description would be an OK way to tackle that: even in your example of a multi-versioned function, all variants would still remain correct, since I only am proposing to add targetMachine attributes. It might be that different variants end up with the same attributes once they've all been added the targetMachine attributes, namely if the targetMachine supports all features required by all variants.... and that would be fine? It's also possible that some variant still differs from others after that, if one of its attribute is outside of what the targetMachine has... and such a variant would be illegal anyway to call on that targetMachine.

[benvanik]

maybe worth a chat? I suspect we may be talking about different things (or I'm just a little lost :) - that our codegen flow today assumes a single target machine with a single set of features is a bug we need to fix, so I'm not sure how we can implement overrideTargetAttributes or what it would give us vs removing the attributes and having the global option (set on the module) take effect. may be worth looking at some bitcode (both generated and from our output) to see what it has and line that up with where we want to be.