From 568aae8667f2704233f763f5e7faeff30c907fe0 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Sun, 17 Jul 2022 14:00:37 +0200
Subject: [PATCH] [dxvk] Use SPIR-V pass to swizzle FS outputs instead of spec
 constants

---
 src/dxvk/dxvk_graphics.cpp | 44 ++++++++++++++++++++------------------
 src/dxvk/dxvk_graphics.h   |  4 ++++
 src/dxvk/dxvk_shader.cpp   |  4 ++++
 3 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/src/dxvk/dxvk_graphics.cpp b/src/dxvk/dxvk_graphics.cpp
index dc8293f59..0d00184f7 100644
--- a/src/dxvk/dxvk_graphics.cpp
+++ b/src/dxvk/dxvk_graphics.cpp
@@ -732,17 +732,8 @@ namespace dxvk {
 
     // Remapping fragment shader outputs would require spec constants
     for (uint32_t i = 0; i < MaxNumRenderTargets; i++) {
-      VkFormat rtFormat = state.rt.getColorFormat(i);
-
-      if (rtFormat && (m_fsOut & (1u << i)) && state.omBlend[i].colorWriteMask()) {
-        auto mapping = state.omSwizzle[i].mapping();
-
-        if (mapping.r != VK_COMPONENT_SWIZZLE_R
-         || mapping.g != VK_COMPONENT_SWIZZLE_G
-         || mapping.b != VK_COMPONENT_SWIZZLE_B
-         || mapping.a != VK_COMPONENT_SWIZZLE_A)
-          return false;
-      }
+      if (writesRenderTarget(state, i) && !util::isIdentityMapping(state.omSwizzle[i].mapping()))
+        return false;
     }
 
     return true;
@@ -813,14 +804,6 @@ namespace dxvk {
     // Set up some specialization constants
     DxvkSpecConstants specData;
 
-    for (uint32_t i = 0; i < MaxNumRenderTargets; i++) {
-      if ((m_fsOut & (1 << i)) != 0) {
-        specData.set(uint32_t(DxvkSpecConstantId::ColorComponentMappings) + i,
-          state.omSwizzle[i].rIndex() << 0 | state.omSwizzle[i].gIndex() << 4 |
-          state.omSwizzle[i].bIndex() << 8 | state.omSwizzle[i].aIndex() << 12, 0x3210u);
-      }
-    }
-
     for (uint32_t i = 0; i < MaxNumSpecConstants; i++)
       specData.set(getSpecId(i), state.sc.specConstants[i], 0u);
     
@@ -907,9 +890,15 @@ namespace dxvk {
     DxvkShaderModuleCreateInfo info;
 
     // Fix up fragment shader outputs for dual-source blending
-    if (shaderInfo.stage == VK_SHADER_STAGE_FRAGMENT_BIT)
+    if (shaderInfo.stage == VK_SHADER_STAGE_FRAGMENT_BIT) {
       info.fsDualSrcBlend = state.useDualSourceBlending();
 
+      for (uint32_t i = 0; i < MaxNumRenderTargets; i++) {
+        if (writesRenderTarget(state, i))
+          info.rtSwizzles[i] = state.omSwizzle[i].mapping();
+      }
+    }
+
     // Deal with undefined shader inputs
     uint32_t consumedInputs = shaderInfo.inputMask;
     uint32_t providedInputs = 0;
@@ -927,7 +916,6 @@ namespace dxvk {
     }
 
     info.undefinedInputs = (providedInputs & consumedInputs) ^ consumedInputs;
-
     return shader->getCode(m_bindings, info);
   }
 
@@ -957,6 +945,20 @@ namespace dxvk {
   }
 
 
+  bool DxvkGraphicsPipeline::writesRenderTarget(
+    const DxvkGraphicsPipelineStateInfo& state,
+          uint32_t                       target) const {
+    if (!(m_fsOut & (1u << target)))
+      return false;
+
+    if (!state.omBlend[target].colorWriteMask())
+      return false;
+
+    VkFormat rtFormat = state.rt.getColorFormat(target);
+    return rtFormat != VK_FORMAT_UNDEFINED;
+  }
+
+
   bool DxvkGraphicsPipeline::validatePipelineState(
     const DxvkGraphicsPipelineStateInfo&  state,
           bool                            trusted) const {
diff --git a/src/dxvk/dxvk_graphics.h b/src/dxvk/dxvk_graphics.h
index 70dbcc1c0..8ab88b746 100644
--- a/src/dxvk/dxvk_graphics.h
+++ b/src/dxvk/dxvk_graphics.h
@@ -432,6 +432,10 @@ namespace dxvk {
     Rc<DxvkShader> getPrevStageShader(
             VkShaderStageFlagBits          stage) const;
 
+    bool writesRenderTarget(
+      const DxvkGraphicsPipelineStateInfo& state,
+            uint32_t                       target) const;
+
     bool validatePipelineState(
       const DxvkGraphicsPipelineStateInfo& state,
             bool                           trusted) const;
diff --git a/src/dxvk/dxvk_shader.cpp b/src/dxvk/dxvk_shader.cpp
index f53dc7ebd..61b57bfd4 100644
--- a/src/dxvk/dxvk_shader.cpp
+++ b/src/dxvk/dxvk_shader.cpp
@@ -140,6 +140,10 @@ namespace dxvk {
     for (uint32_t u : bit::BitMask(state.undefinedInputs))
       eliminateInput(spirvCode, u);
 
+    // Emit fragment shader swizzles as necessary
+    if (m_info.stage == VK_SHADER_STAGE_FRAGMENT_BIT)
+      emitOutputSwizzles(spirvCode, m_info.outputMask, state.rtSwizzles.data());
+
     return spirvCode;
   }