[AMDGPU] Add AMDGPU-specific module splitting

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

@@ -20,6 +20,7 @@
 #include "AMDGPUIGroupLP.h"
 #include "AMDGPUMacroFusion.h"
 #include "AMDGPURegBankSelect.h"
+#include "AMDGPUSplitModule.h"
 #include "AMDGPUTargetObjectFile.h"
 #include "AMDGPUTargetTransformInfo.h"
 #include "AMDGPUUnifyDivergentExitNodes.h"
@@ -806,6 +807,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

@@ -67,6 +67,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

@@ -97,6 +97,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 -mtarget 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 -mtarget 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,54 @@
+; RUN: llvm-split -o %t %s -j 3 -mtarget 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() {
+    store i64 42, ptr undef
+    store i64 43, ptr undef
+    store i64 44, ptr undef
+    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
+}

llvm/test/tools/llvm-split/AMDGPU/kernels-cost-ranking.ll

@@ -0,0 +1,50 @@
+; RUN: llvm-split -o %t %s -j 3 -mtarget 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.
+
+; 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 @B
+; CHECK1: define internal void @HelperB
+; CHECK1-NOT: define
+
+; CHECK2-NOT: define
+; CHECK2: define amdgpu_kernel void @A
+; CHECK2: define internal void @HelperA
+; CHECK2-NOT: define
+
+
+define amdgpu_kernel void @A() {
+    call void @HelperA()
+    ret void
+}
+
+define internal void @HelperA() {
+    ret void
+}
+
+define amdgpu_kernel void @B() {
+    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
+}

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

@@ -0,0 +1,41 @@
+; RUN: llvm-split -o %t %s -j 3 -mtarget 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 share a common helper, that helper should be
+; cloned in all partitions.
+
+; CHECK0-NOT: define
+; CHECK0: define internal void @Helper
+; CHECK0: define amdgpu_kernel void @C
+; CHECK0-NOT: define
+
+; CHECK1-NOT: define
+; CHECK1: define internal void @Helper
+; CHECK1: define amdgpu_kernel void @B
+; CHECK1-NOT: define
+
+; CHECK2-NOT: define
+; CHECK2: define internal void @Helper
+; CHECK2: define amdgpu_kernel void @A
+; CHECK2-NOT: define
+
+define internal void @Helper() {
+    ret void
+}
+
+define amdgpu_kernel void @A() {
+    call void @Helper()
+    ret void
+}
+
+define amdgpu_kernel void @B() {
+    call void @Helper()
+    ret void
+}
+
+define amdgpu_kernel void @C() {
+    call void @Helper()
+    ret void
+}

llvm/test/tools/llvm-split/AMDGPU/kernels-dependency-duplication.ll

@@ -0,0 +1,43 @@
+; RUN: llvm-split -o %t %s -j 3 -mtarget 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 use private/internal global variables.
+; The GVs should be copied in each partition as needed.
+
+; CHECK0-NOT: define
+; CHECK0: define void @ExternalHelper
+; CHECK0: define amdgpu_kernel void @A
+; CHECK0: define amdgpu_kernel void @B
+; CHECK0-NOT: define
+
+; CHECK1-NOT: define
+; CHECK1: define amdgpu_kernel void @D
+; CHECK1-NOT: define
+
+; CHECK2-NOT: define
+; CHECK2: define amdgpu_kernel void @C
+; CHECK2-NOT: define
+
+define void @ExternalHelper() {
+    ret void
+}
+
+define amdgpu_kernel void @A() {
+    call void @ExternalHelper()
+    ret void
+}
+
+define amdgpu_kernel void @B() {
+    call void @ExternalHelper()
+    ret void
+}
+
+define amdgpu_kernel void @C() {
+    ret void
+}
+
+define amdgpu_kernel void @D() {
+    ret void
+}