Index: llvm/include/llvm/ADT/GenericUniformityImpl.h
===================================================================
--- llvm/include/llvm/ADT/GenericUniformityImpl.h
+++ llvm/include/llvm/ADT/GenericUniformityImpl.h
@@ -43,6 +43,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SparseBitVector.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <set>
@@ -1140,7 +1141,6 @@
     : F(&Func) {
   DA.reset(new ImplT{Func, DT, CI, TTI});
   DA->initialize();
-  DA->compute();
 }
 
 template <typename ContextT>
Index: llvm/include/llvm/ADT/GenericUniformityInfo.h
===================================================================
--- llvm/include/llvm/ADT/GenericUniformityInfo.h
+++ llvm/include/llvm/ADT/GenericUniformityInfo.h
@@ -47,6 +47,10 @@
   GenericUniformityInfo(GenericUniformityInfo &&) = default;
   GenericUniformityInfo &operator=(GenericUniformityInfo &&) = default;
 
+  void compute() {
+    DA->compute();
+  }
+
   /// Whether any divergence was detected.
   bool hasDivergence() const;
 
Index: llvm/include/llvm/CodeGen/MachineUniformityAnalysis.h
===================================================================
--- llvm/include/llvm/CodeGen/MachineUniformityAnalysis.h
+++ llvm/include/llvm/CodeGen/MachineUniformityAnalysis.h
@@ -25,10 +25,12 @@
 using MachineUniformityInfo = GenericUniformityInfo<MachineSSAContext>;
 
 /// \brief Compute uniformity information for a Machine IR function.
-MachineUniformityInfo
-computeMachineUniformityInfo(MachineFunction &F,
-                             const MachineCycleInfo &cycleInfo,
-                             const MachineDomTree &domTree);
+///
+/// If \p HasBranchDivergence is false, produces a dummy result which assumes
+/// everything is uniform.
+MachineUniformityInfo computeMachineUniformityInfo(
+    MachineFunction &F, const MachineCycleInfo &cycleInfo,
+    const MachineDomTree &domTree, bool HasBranchDivergence);
 
 } // namespace llvm
 
Index: llvm/lib/Analysis/UniformityAnalysis.cpp
===================================================================
--- llvm/lib/Analysis/UniformityAnalysis.cpp
+++ llvm/lib/Analysis/UniformityAnalysis.cpp
@@ -168,6 +168,11 @@
   m_function = &F;
   m_uniformityInfo =
       UniformityInfo{F, domTree, cycleInfo, &targetTransformInfo};
+
+  // Skip computation if we can assume everything is uniform.
+  if (targetTransformInfo.hasBranchDivergence())
+    m_uniformityInfo.compute();
+
   return false;
 }
 
Index: llvm/lib/CodeGen/MachineUniformityAnalysis.cpp
===================================================================
--- llvm/lib/CodeGen/MachineUniformityAnalysis.cpp
+++ llvm/lib/CodeGen/MachineUniformityAnalysis.cpp
@@ -153,12 +153,14 @@
 template struct llvm::GenericUniformityAnalysisImplDeleter<
     llvm::GenericUniformityAnalysisImpl<MachineSSAContext>>;
 
-MachineUniformityInfo
-llvm::computeMachineUniformityInfo(MachineFunction &F,
-                                   const MachineCycleInfo &cycleInfo,
-                                   const MachineDomTree &domTree) {
+MachineUniformityInfo llvm::computeMachineUniformityInfo(
+    MachineFunction &F, const MachineCycleInfo &cycleInfo,
+    const MachineDomTree &domTree, bool HasBranchDivergence) {
   assert(F.getRegInfo().isSSA() && "Expected to be run on SSA form!");
-  return MachineUniformityInfo(F, domTree, cycleInfo);
+  MachineUniformityInfo UI(F, domTree, cycleInfo);
+  if (HasBranchDivergence)
+    UI.compute();
+  return UI;
 }
 
 namespace {
@@ -205,6 +207,7 @@
                       "Machine Uniformity Info Analysis", true, true)
 INITIALIZE_PASS_DEPENDENCY(MachineCycleInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(MachineUniformityAnalysisPass, "machine-uniformity",
                     "Machine Uniformity Info Analysis", true, true)
 
@@ -218,7 +221,9 @@
 bool MachineUniformityAnalysisPass::runOnMachineFunction(MachineFunction &MF) {
   auto &DomTree = getAnalysis<MachineDominatorTree>().getBase();
   auto &CI = getAnalysis<MachineCycleInfoWrapperPass>().getCycleInfo();
-  UI = computeMachineUniformityInfo(MF, CI, DomTree);
+  auto &TTI =
+      getAnalysis<TargetTransformInfoWrapperPass>().getTTI(MF.getFunction());
+  UI = computeMachineUniformityInfo(MF, CI, DomTree, TTI.hasBranchDivergence());
   return false;
 }
 
Index: llvm/lib/Target/AMDGPU/AMDGPURegBankSelect.cpp
===================================================================
--- llvm/lib/Target/AMDGPU/AMDGPURegBankSelect.cpp
+++ llvm/lib/Target/AMDGPU/AMDGPURegBankSelect.cpp
@@ -63,8 +63,9 @@
       getAnalysis<MachineCycleInfoWrapperPass>().getCycleInfo();
   MachineDominatorTree &DomTree = getAnalysis<MachineDominatorTree>();
 
+  // TODO: Check for single lane execution.
   MachineUniformityInfo Uniformity =
-      computeMachineUniformityInfo(MF, CycleInfo, DomTree.getBase());
+      computeMachineUniformityInfo(MF, CycleInfo, DomTree.getBase(), true);
   (void)Uniformity; // TODO: Use this
 
   assignRegisterBanks(MF);