From accad56ee7cc9f77886d164701a35f1e89a3519b Mon Sep 17 00:00:00 2001
From: ReinUsesLisp <reinuseslisp@airmail.cc>
Date: Wed, 19 May 2021 16:32:03 -0300
Subject: [PATCH] glasm: Implement stores to gl_ViewportIndex

---
 .../backend/glasm/emit_context.cpp               |  4 +++-
 .../backend/glasm/emit_context.h                 |  8 ++++++--
 .../backend/glasm/emit_glasm.cpp                 | 16 ++++++++++++----
 .../backend/glasm/emit_glasm_context_get_set.cpp |  8 ++++++++
 4 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/shader_recompiler/backend/glasm/emit_context.cpp b/src/shader_recompiler/backend/glasm/emit_context.cpp
index e2182400ca..395ac87f20 100644
--- a/src/shader_recompiler/backend/glasm/emit_context.cpp
+++ b/src/shader_recompiler/backend/glasm/emit_context.cpp
@@ -23,7 +23,8 @@ std::string_view InterpDecorator(Interpolation interp) {
 }
 } // Anonymous namespace
 
-EmitContext::EmitContext(IR::Program& program, Bindings& bindings) : info{program.info} {
+EmitContext::EmitContext(IR::Program& program, Bindings& bindings, const Profile& profile_)
+    : info{program.info}, profile{profile_} {
     // FIXME: Temporary partial implementation
     u32 cbuf_index{};
     for (const auto& desc : program.info.constant_buffer_descriptors) {
@@ -41,6 +42,7 @@ EmitContext::EmitContext(IR::Program& program, Bindings& bindings) : info{progra
     if (const size_t num = program.info.storage_buffers_descriptors.size(); num > 0) {
         Add("PARAM c[{}]={{program.local[0..{}]}};", num, num - 1);
     }
+    stage = program.stage;
     switch (program.stage) {
     case Stage::VertexA:
     case Stage::VertexB:
diff --git a/src/shader_recompiler/backend/glasm/emit_context.h b/src/shader_recompiler/backend/glasm/emit_context.h
index d6b0bf73c9..dd1f9ac9f0 100644
--- a/src/shader_recompiler/backend/glasm/emit_context.h
+++ b/src/shader_recompiler/backend/glasm/emit_context.h
@@ -11,10 +11,12 @@
 #include <fmt/format.h>
 
 #include "shader_recompiler/backend/glasm/reg_alloc.h"
+#include "shader_recompiler/stage.h"
 
 namespace Shader {
 struct Info;
-}
+struct Profile;
+} // namespace Shader
 
 namespace Shader::Backend {
 struct Bindings;
@@ -29,7 +31,7 @@ namespace Shader::Backend::GLASM {
 
 class EmitContext {
 public:
-    explicit EmitContext(IR::Program& program, Bindings& bindings);
+    explicit EmitContext(IR::Program& program, Bindings& bindings, const Profile& profile_);
 
     template <typename... Args>
     void Add(const char* format_str, IR::Inst& inst, Args&&... args) {
@@ -55,10 +57,12 @@ public:
     std::string code;
     RegAlloc reg_alloc{*this};
     const Info& info;
+    const Profile& profile;
 
     std::vector<u32> texture_buffer_bindings;
     std::vector<u32> texture_bindings;
 
+    Stage stage{};
     std::string_view stage_name = "invalid";
 };
 
diff --git a/src/shader_recompiler/backend/glasm/emit_glasm.cpp b/src/shader_recompiler/backend/glasm/emit_glasm.cpp
index edf6f5e13e..9dc0cacbea 100644
--- a/src/shader_recompiler/backend/glasm/emit_glasm.cpp
+++ b/src/shader_recompiler/backend/glasm/emit_glasm.cpp
@@ -261,7 +261,10 @@ void EmitCode(EmitContext& ctx, const IR::Program& program) {
     }
 }
 
-void SetupOptions(std::string& header, Info info) {
+void SetupOptions(const IR::Program& program, const Profile& profile, std::string& header) {
+    const Info& info{program.info};
+    const Stage stage{program.stage};
+
     // TODO: Track the shared atomic ops
     header += "OPTION NV_internal;"
               "OPTION NV_shader_storage_buffer;"
@@ -286,6 +289,11 @@ void SetupOptions(std::string& header, Info info) {
     if (info.uses_sparse_residency) {
         header += "OPTION EXT_sparse_texture2;";
     }
+    if ((info.stores_viewport_index || info.stores_layer) && stage != Stage::Geometry) {
+        if (profile.support_viewport_index_layer_non_geometry) {
+            header += "OPTION NV_viewport_array2;";
+        }
+    }
     const auto non_zero_frag_colors{info.stores_frag_color | std::views::drop(1)};
     if (std::ranges::find(non_zero_frag_colors, true) != non_zero_frag_colors.end()) {
         header += "OPTION ARB_draw_buffers;";
@@ -312,12 +320,12 @@ std::string_view StageHeader(Stage stage) {
 }
 } // Anonymous namespace
 
-std::string EmitGLASM(const Profile&, IR::Program& program, Bindings& bindings) {
-    EmitContext ctx{program, bindings};
+std::string EmitGLASM(const Profile& profile, IR::Program& program, Bindings& bindings) {
+    EmitContext ctx{program, bindings, profile};
     Precolor(ctx, program);
     EmitCode(ctx, program);
     std::string header{StageHeader(program.stage)};
-    SetupOptions(header, program.info);
+    SetupOptions(program, profile, header);
     switch (program.stage) {
     case Stage::Compute:
         header += fmt::format("GROUP_SIZE {} {} {};", program.workgroup_size[0],
diff --git a/src/shader_recompiler/backend/glasm/emit_glasm_context_get_set.cpp b/src/shader_recompiler/backend/glasm/emit_glasm_context_get_set.cpp
index f362dd2c8f..6484387bcb 100644
--- a/src/shader_recompiler/backend/glasm/emit_glasm_context_get_set.cpp
+++ b/src/shader_recompiler/backend/glasm/emit_glasm_context_get_set.cpp
@@ -7,6 +7,7 @@
 #include "shader_recompiler/backend/glasm/emit_context.h"
 #include "shader_recompiler/backend/glasm/emit_glasm_instructions.h"
 #include "shader_recompiler/frontend/ir/value.h"
+#include "shader_recompiler/profile.h"
 
 namespace Shader::Backend::GLASM {
 namespace {
@@ -102,6 +103,13 @@ void EmitSetAttribute(EmitContext& ctx, IR::Attribute attr, ScalarF32 value,
     case IR::Attribute::PositionW:
         ctx.Add("MOV.F result.position.{},{};", swizzle, value);
         break;
+    case IR::Attribute::ViewportIndex:
+        if (ctx.stage == Stage::Geometry || ctx.profile.support_viewport_index_layer_non_geometry) {
+            ctx.Add("MOV.F result.viewport.x,{};", value);
+        } else {
+            // LOG_WARNING
+        }
+        break;
     default:
         throw NotImplementedException("Set attribute {}", attr);
     }