[AMDGPU] Add AMDGPU-specific module splitting (#89245)

This enables the --lto-partitions option to work more consistently.

This module splitting logic is fully aware of AMDGPU modules and their
specificities and takes advantage of
them to split modules in a way that avoids compilation issue (such as
resource usage being incorrectly represented).

This also includes a logging system that's more elaborate than just
LLVM_DEBUG which allows
printing logs to uniquely named files, and optionally with all value
names hidden so they can be safely shared without leaking informatiton
about the source. Logs can also be enabled through an environment
variable, which avoids the sometimes complicated process of passing a
-mllvm option all the way from clang driver to the offload linker that
handles full LTO codegen.

Author: Pierre-vh

llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp

@@ -0,0 +1,30 @@
+//===- AMDGPUSplitModule.h -------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TARGET_AMDGPUSPLITMODULE_H
+#define LLVM_TARGET_AMDGPUSPLITMODULE_H
+
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include <memory>
+
+namespace llvm {
+
+class Module;
+class AMDGPUTargetMachine;
+
+/// Splits the module M into N linkable partitions. The function ModuleCallback
+/// is called N times passing each individual partition as the MPart argument.
+void splitAMDGPUModule(
+    const AMDGPUTargetMachine &TM, Module &M, unsigned N,
+    function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback);
+
+} // end namespace llvm
+
+#endif // LLVM_TARGET_AMDGPUSPLITMODULE_H

llvm/lib/Target/AMDGPU/AMDGPUSplitModule.h

@@ -21,6 +21,7 @@
 #include "AMDGPUIGroupLP.h"
 #include "AMDGPUMacroFusion.h"
 #include "AMDGPURegBankSelect.h"
+#include "AMDGPUSplitModule.h"
 #include "AMDGPUTargetObjectFile.h"
 #include "AMDGPUTargetTransformInfo.h"
 #include "AMDGPUUnifyDivergentExitNodes.h"
@@ -815,6 +816,13 @@ AMDGPUTargetMachine::getAddressSpaceForPseudoSourceKind(unsigned Kind) const {
   return AMDGPUAS::FLAT_ADDRESS;
 }
 
+bool AMDGPUTargetMachine::splitModule(
+    Module &M, unsigned NumParts,
+    function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback) const {
+  splitAMDGPUModule(*this, M, NumParts, ModuleCallback);
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // GCN Target Machine (SI+)
 //===----------------------------------------------------------------------===//

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -73,6 +73,10 @@ class AMDGPUTargetMachine : public LLVMTargetMachine {
   getPredicatedAddrSpace(const Value *V) const override;
 
   unsigned getAddressSpaceForPseudoSourceKind(unsigned Kind) const override;
+
+  bool splitModule(Module &M, unsigned NumParts,
+                   function_ref<void(std::unique_ptr<Module> MPart)>
+                       ModuleCallback) const override;
 };
 
 //===----------------------------------------------------------------------===//

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h

@@ -98,6 +98,7 @@ add_llvm_target(AMDGPUCodeGen
   AMDGPURewriteOutArguments.cpp
   AMDGPURewriteUndefForPHI.cpp
   AMDGPUSetWavePriority.cpp
+  AMDGPUSplitModule.cpp
   AMDGPUSubtarget.cpp
   AMDGPUTargetMachine.cpp
   AMDGPUTargetObjectFile.cpp

llvm/lib/Target/AMDGPU/CMakeLists.txt

@@ -0,0 +1,46 @@
+; RUN: llvm-split -o %t %s -j 3 -mtriple amdgcn-amd-amdhsa -amdgpu-module-splitting-large-kernel-threshold=0
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t2 | FileCheck --check-prefix=CHECK2 %s
+
+; 3 kernels:
+;   - A does a direct call to HelperA
+;   - B is storing @HelperA
+;   - C does a direct call to HelperA
+;
+; The helper functions will get externalized, which will force A and C into P0 as
+; external functions cannot be duplicated.
+
+; CHECK0: define hidden void @HelperA()
+; CHECK0: define amdgpu_kernel void @A()
+; CHECK0: declare amdgpu_kernel void @B(ptr)
+; CHECK0: define amdgpu_kernel void @C()
+
+; CHECK1: declare hidden void @HelperA()
+; CHECK1: declare amdgpu_kernel void @A()
+; CHECK1: declare amdgpu_kernel void @B(ptr)
+; CHECK1: declare amdgpu_kernel void @C()
+
+; CHECK2: declare hidden void @HelperA()
+; CHECK2: declare amdgpu_kernel void @A()
+; CHECK2: define amdgpu_kernel void @B(ptr %dst)
+; CHECK2: declare amdgpu_kernel void @C()
+
+define internal void @HelperA() {
+  ret void
+}
+
+define amdgpu_kernel void @A() {
+  call void @HelperA()
+  ret void
+}
+
+define amdgpu_kernel void @B(ptr %dst) {
+  store ptr @HelperA, ptr %dst
+  ret void
+}
+
+define amdgpu_kernel void @C() {
+  call void @HelperA()
+  ret void
+}

llvm/test/tools/llvm-split/AMDGPU/address-taken-externalize-with-call.ll

@@ -0,0 +1,37 @@
+; RUN: llvm-split -o %t %s -j 2 -mtriple amdgcn-amd-amdhsa -amdgpu-module-splitting-large-kernel-threshold=0
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+
+; 2 kernels:
+;   - A is isolated
+;   - B is storing @HelperA/B's address
+;
+; The helper functions should get externalized (become hidden w/ external linkage)
+
+; CHECK0: define hidden void @HelperA()
+; CHECK0: define hidden void @HelperB()
+; CHECK0: define amdgpu_kernel void @A()
+; CHECK0: declare amdgpu_kernel void @B(i1, ptr)
+
+; CHECK1: declare hidden void @HelperA()
+; CHECK1: declare hidden void @HelperB()
+; CHECK1: declare amdgpu_kernel void @A()
+; CHECK1: define amdgpu_kernel void @B(i1 %cond, ptr %dst)
+
+define internal void @HelperA() {
+  ret void
+}
+
+define internal void @HelperB() {
+  ret void
+}
+
+define amdgpu_kernel void @A() {
+  ret void
+}
+
+define amdgpu_kernel void @B(i1 %cond, ptr %dst) {
+  %addr = select i1 %cond, ptr @HelperA, ptr @HelperB
+  store ptr %addr, ptr %dst
+  ret void
+}

llvm/test/tools/llvm-split/AMDGPU/address-taken-externalize.ll

@@ -0,0 +1,20 @@
+; RUN: llvm-split -o %t %s -j 3 -mtriple amdgcn-amd-amdhsa -debug -amdgpu-module-splitting-log-private 2>&1 | FileCheck %s --implicit-check-not=MyCustomKernel
+; REQUIRES: asserts
+
+; SHA256 of the kernel names.
+
+; CHECK: a097723d21cf9f35d90e6fb7881995ac8c398b3366a6c97efc657404f9fe301c
+; CHECK: 626bc23242de8fcfda7f0e66318d29455c081df6b5380e64d14703c95fcbcd59
+; CHECK: c38d90a7ca71dc5d694bb9e093dadcdedfc4cb4adf7ed7e46d42fe95a0b4ef55
+
+define amdgpu_kernel void @MyCustomKernel0() {
+  ret void
+}
+
+define amdgpu_kernel void @MyCustomKernel1() {
+  ret void
+}
+
+define amdgpu_kernel void @MyCustomKernel2() {
+  ret void
+}

llvm/test/tools/llvm-split/AMDGPU/debug-name-hiding.ll

@@ -0,0 +1,45 @@
+; RUN: llvm-split -o %t %s -j 2 -mtriple amdgcn-amd-amdhsa
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+
+; 3 kernels:
+;   - A calls nothing
+;   - B calls @PerryThePlatypus
+;   - C calls @Perry, an alias of @PerryThePlatypus
+;
+; We should see through the alias and put B/C in the same
+; partition.
+;
+; Additionally, @PerryThePlatypus gets externalized as
+; the alias counts as taking its address.
+
+; CHECK0-NOT: define
+; CHECK0: @Perry = internal alias ptr (), ptr @PerryThePlatypus
+; CHECK0: define hidden void @PerryThePlatypus()
+; CHECK0: define amdgpu_kernel void @B
+; CHECK0: define amdgpu_kernel void @C
+; CHECK0-NOT: define
+
+; CHECK1-NOT: define
+; CHECK1: define amdgpu_kernel void @A
+; CHECK1-NOT: define
+
+@Perry = internal alias ptr(), ptr @PerryThePlatypus
+
+define internal void @PerryThePlatypus() {
+  ret void
+}
+
+define amdgpu_kernel void @A() {
+  ret void
+}
+
+define amdgpu_kernel void @B() {
+  call void @PerryThePlatypus()
+  ret void
+}
+
+define amdgpu_kernel void @C() {
+  call void @Perry()
+  ret void
+}

llvm/test/tools/llvm-split/AMDGPU/kernels-alias-dependencies.ll

@@ -0,0 +1,54 @@
+; RUN: llvm-split -o %t %s -j 3 -mtriple amdgcn-amd-amdhsa
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t2 | FileCheck --check-prefix=CHECK2 %s
+
+; 3 kernels with each their own dependencies should go into 3
+; distinct partitions. The most expensive kernel should be
+; seen first and go into the last partition.
+
+; CHECK0-NOT: define
+; CHECK0: define amdgpu_kernel void @C
+; CHECK0: define internal void @HelperC
+; CHECK0-NOT: define
+
+; CHECK1-NOT: define
+; CHECK1: define amdgpu_kernel void @A
+; CHECK1: define internal void @HelperA
+; CHECK1-NOT: define
+
+; CHECK2-NOT: define
+; CHECK2: define amdgpu_kernel void @B
+; CHECK2: define internal void @HelperB
+; CHECK2-NOT: define
+
+
+define amdgpu_kernel void @A() {
+  call void @HelperA()
+  ret void
+}
+
+define internal void @HelperA() {
+  ret void
+}
+
+define amdgpu_kernel void @B(ptr %x) {
+  store i64 42, ptr %x
+  store i64 43, ptr %x
+  store i64 44, ptr %x
+  call void @HelperB()
+  ret void
+}
+
+define internal void @HelperB() {
+  ret void
+}
+
+define amdgpu_kernel void @C() {
+  call void @HelperC()
+  ret void
+}
+
+define internal void @HelperC() {
+  ret void
+}