[AMDGPU] Cluster shader exports

[critson]

Summary

Add DAG scheduling mutation to cluster export instructions.
This avoids unnecessary waitcnts being added when computation
ends up interspersed with exports.

[critson]

Apply linting.

[foad]

Looks good.

As a follow-on, would this be a good place to try to ensure that position exports go before parameter exports?

llvm/lib/Target/AMDGPU/AMDGPUExportClustering.cpp

32–33

You don't need these parentheses.

40

This is trivia, but isn't it simpler to start counting at 1 than to finish at (End - 1) ?

46

I don't quite understand the comment. If nothing depends on an export then surely isExport() on the next line will never be true?

59–79

Do you understand the chain formation logic, or is it copied from the standard MemOpClusterMutation? (Or both?!)

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.exp.ll

552–558

Can you pre-commit this? Did it actually have a waitcnt before this fix?

[critson]

llvm/lib/Target/AMDGPU/AMDGPUExportClustering.cpp

59–79

A bit of both, code adapted from MemOpClusterMutation and I have reasonable confidence I understand what it is doing.

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.exp.ll

552–558

Sorry for the effort, but can you educate me on the concept of pre-committing the test. It will fail without the rest of the code in this patch as a waitcnt is inserted at present.

[critson]

Address Jay's comments.

[foad]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.exp.ll

552–558

I mean, commit @test_export_clustering but with check lines that show the current codegen. (This is easier if you auto-generate the check lines. I hadn't noticed that you wrote these ones by hand.) Then, when you rebase the current patch, it will show us exactly what effect it had on the codegen for this function.

It's not essential but it is helpful to reviewers.

[critson]

Precommit test.

[foad]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.exp.ll

552–558

Nice!

llvm/lib/Target/AMDGPU/AMDGPUExportClustering.h

+//===- AMDGPUExportClustering.h - AMDGPU Export Clustering ------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MachineScheduler.h"
+
+namespace llvm {
+
+std::unique_ptr<ScheduleDAGMutation> createAMDGPUExportClusteringDAGMutation();
+
+} // namespace llvm

llvm/lib/Target/AMDGPU/AMDGPUExportClustering.cpp

+//===--- AMDGPUExportClusting.cpp - AMDGPU Export Clustering  -------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file This file contains a DAG scheduling mutation to cluster shader
+///       exports.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPUExportClustering.h"
+#include "AMDGPUSubtarget.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIInstrInfo.h"
+
+using namespace llvm;
+
+namespace {
+
+class ExportClustering : public ScheduleDAGMutation {
+public:
+  ExportClustering() {}
+  void apply(ScheduleDAGInstrs *DAG) override;
+};
+
+static bool isExport(const SUnit &SU) {
+  const MachineInstr *MI = SU.getInstr();
+  return MI->getOpcode() == AMDGPU::EXP ||
+         MI->getOpcode() == AMDGPU::EXP_DONE;
+}
+
+static void buildCluster(ArrayRef<SUnit *> Exports, ScheduleDAGInstrs *DAG) {
+  // Cluster a series of exports. Also copy all dependencies to the first
+  // export to avoid computation being inserted into the chain.
+  SUnit *ChainHead = Exports[0];
+  for (unsigned Idx = 0, End = Exports.size() - 1; Idx < End; ++Idx) {
+    SUnit *SUa = Exports[Idx];
+    SUnit *SUb = Exports[Idx + 1];
+    if (DAG->addEdge(SUb, SDep(SUa, SDep::Cluster))) {
+      for (const SDep &Pred : SUb->Preds) {
+        SUnit *PredSU = Pred.getSUnit();
+        if (Pred.isWeak() || isExport(*PredSU))
+          continue;
+        DAG->addEdge(ChainHead, SDep(PredSU, SDep::Artificial));
+      }
+    }
+  }
+}
+
+void ExportClustering::apply(ScheduleDAGInstrs *DAG) {
+  SmallVector<SmallVector<SUnit *, 8>, 4> ExportChains;
+  DenseMap<unsigned, unsigned> ChainMap;
+
+  // Build chains of exports
+  for (SUnit &SU : DAG->SUnits) {
+    if (!isExport(SU))
+      continue;
+
+    unsigned ChainID = ExportChains.size();
+    for (const SDep &Pred : SU.Preds) {
+      const SUnit &PredSU = *Pred.getSUnit();
+      if (isExport(PredSU) && !Pred.isArtificial()) {
+        ChainID = ChainMap.lookup(PredSU.NodeNum);
+        break;
+      }
+    }
+    ChainMap[SU.NodeNum] = ChainID;
+
+    if (ChainID == ExportChains.size())
+      ExportChains.push_back(SmallVector<SUnit *, 8>());
+
+    auto &Chain = ExportChains[ChainID];
+    Chain.push_back(&SU);
+  }
+
+  // Apply clustering
+  for (auto &Chain : ExportChains)
+    buildCluster(Chain, DAG);
+}
+
+} // end namespace
+
+namespace llvm {
+
+std::unique_ptr<ScheduleDAGMutation> createAMDGPUExportClusteringDAGMutation() {
+  return std::make_unique<ExportClustering>();
+}
+
+} // end namespace llvm

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

 //===-- AMDGPUTargetMachine.cpp - TargetMachine for hw codegen targets-----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 /// The AMDGPU target machine contains all of the hardware specific
 /// information  needed to emit code for R600 and SI GPUs.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPUTargetMachine.h"
 #include "AMDGPU.h"
 #include "AMDGPUAliasAnalysis.h"
 #include "AMDGPUCallLowering.h"
+#include "AMDGPUExportClustering.h"
 #include "AMDGPUInstructionSelector.h"
 #include "AMDGPULegalizerInfo.h"
 #include "AMDGPUMacroFusion.h"
 #include "AMDGPUTargetObjectFile.h"
 #include "AMDGPUTargetTransformInfo.h"
 #include "GCNIterativeScheduler.h"
 #include "GCNSchedStrategy.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"

 
 static ScheduleDAGInstrs *
 createGCNMaxOccupancyMachineScheduler(MachineSchedContext *C) {
   ScheduleDAGMILive *DAG =
     new GCNScheduleDAGMILive(C, std::make_unique<GCNMaxOccupancySchedStrategy>(C));
   DAG->addMutation(createLoadClusterDAGMutation(DAG->TII, DAG->TRI));
   DAG->addMutation(createStoreClusterDAGMutation(DAG->TII, DAG->TRI));
   DAG->addMutation(createAMDGPUMacroFusionDAGMutation());
+  DAG->addMutation(createAMDGPUExportClusteringDAGMutation());
   return DAG;
 }
 
 static ScheduleDAGInstrs *
 createIterativeGCNMaxOccupancyMachineScheduler(MachineSchedContext *C) {
   auto DAG = new GCNIterativeScheduler(C,
     GCNIterativeScheduler::SCHEDULE_LEGACYMAXOCCUPANCY);
   DAG->addMutation(createLoadClusterDAGMutation(DAG->TII, DAG->TRI));

llvm/lib/Target/AMDGPU/CMakeLists.txt

add_llvm_target(AMDGPUCodeGen
   AMDGPUAlwaysInlinePass.cpp
   AMDGPUAnnotateKernelFeatures.cpp
   AMDGPUAnnotateUniformValues.cpp
   AMDGPUArgumentUsageInfo.cpp
   AMDGPUAsmPrinter.cpp
   AMDGPUAtomicOptimizer.cpp
   AMDGPUCallLowering.cpp
   AMDGPUCodeGenPrepare.cpp
+  AMDGPUExportClustering.cpp
   AMDGPUFixFunctionBitcasts.cpp
   AMDGPUFrameLowering.cpp
   AMDGPUHSAMetadataStreamer.cpp
   AMDGPUInstrInfo.cpp
   AMDGPUInstructionSelector.cpp
   AMDGPUISelDAGToDAG.cpp
   AMDGPUISelLowering.cpp
   AMDGPUGlobalISelUtils.cpp

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.exp.ll

exp:
   br label %end
 
 end:
   ret void
 }
 
 ; GCN-LABEL: {{^}}test_export_clustering:
 ; GCN-DAG: v_mov_b32_e32 [[W0:v[0-9]+]], 0
+; GCN-DAG: v_mov_b32_e32 [[W1:v[0-9]+]], 1.0
 ; GCN-DAG: v_mov_b32_e32 [[X:v[0-9]+]], s0
 ; GCN-DAG: v_mov_b32_e32 [[Y:v[0-9]+]], s1
 ; GCN-DAG: v_add_f32_e32 [[Z0:v[0-9]+]]
-; GCN-DAG: exp param0 [[X]], [[Y]], [[Z0]], [[W0]]{{$}}
 ; GCN-DAG: v_sub_f32_e32 [[Z1:v[0-9]+]]
-; GCN: s_waitcnt expcnt(0)
-; GCN: v_mov_b32_e32 [[W1:v[0-9]+]], 1.0
-; GCN: exp param1 [[X]], [[Y]], [[Z1]], [[W1]] done{{$}}
+; GCN: exp param0 [[X]], [[Y]], [[Z0]], [[W0]]{{$}}
+; GCN-NEXT: exp param1 [[X]], [[Y]], [[Z1]], [[W1]] done{{$}}
 define amdgpu_kernel void @test_export_clustering(float %x, float %y) #0 {
   %z0 = fadd float %x, %y
   call void @llvm.amdgcn.exp.f32(i32 32, i32 15, float %x, float %y, float %z0, float 0.0, i1 false, i1 false)
   %z1 = fsub float %y, %x
   call void @llvm.amdgcn.exp.f32(i32 33, i32 15, float %x, float %y, float %z1, float 1.0, i1 true, i1 false)
   ret void
 }
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind inaccessiblememonly }