7 files changed, 32 insertions, 2 deletions

diff --git a/src/shader_recompiler/backend/spirv/emit_context.cpp b/src/shader_recompiler/backend/spirv/emit_context.cpp
index 3b3fea50c..a8041aadc 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_context.cpp
@@ -509,6 +509,12 @@ void EmitContext::DefineOutputs(const Info& info) {
         const Id type{TypeArray(F32[1], Constant(U32[1], 8U))};
         clip_distances = DefineOutput(*this, type, spv::BuiltIn::ClipDistance);
     }
+    if (info.stores_viewport_index && !ignore_viewport_layer) {
+        if (stage == Stage::Fragment) {
+            throw NotImplementedException("Storing ViewportIndex in Fragment stage");
+        }
+        viewport_index = DefineOutput(*this, U32[1], spv::BuiltIn::ViewportIndex);
+    }
     for (size_t i = 0; i < info.stores_generics.size(); ++i) {
         if (info.stores_generics[i]) {
             output_generics[i] = DefineOutput(*this, F32[4]);
diff --git a/src/shader_recompiler/backend/spirv/emit_context.h b/src/shader_recompiler/backend/spirv/emit_context.h
index dc4e1227a..1573c2560 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.h
+++ b/src/shader_recompiler/backend/spirv/emit_context.h
@@ -106,6 +106,7 @@ public:
     Id front_face{};
     Id point_coord{};
     Id clip_distances{};
+    Id viewport_index{};
 
     Id fswzadd_lut_a{};
     Id fswzadd_lut_b{};
@@ -133,6 +134,8 @@ public:
 
     std::vector<Id> interfaces;
 
+    bool ignore_viewport_layer{};
+
 private:
     void DefineCommonTypes(const Info& info);
     void DefineCommonConstants();
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index b8e3b8527..cc6b98f7e 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -225,6 +225,16 @@ void SetupCapabilities(const Profile& profile, const Info& info, EmitContext& ct
         ctx.AddExtension("SPV_EXT_demote_to_helper_invocation");
         ctx.AddCapability(spv::Capability::DemoteToHelperInvocationEXT);
     }
+    if (info.stores_viewport_index) {
+        ctx.AddCapability(spv::Capability::MultiViewport);
+        if (profile.support_viewport_index_layer_non_geometry &&
+            ctx.stage == Shader::Stage::VertexB) {
+            ctx.AddExtension("SPV_EXT_shader_viewport_index_layer");
+            ctx.AddCapability(spv::Capability::ShaderViewportIndexLayerEXT);
+        } else {
+            ctx.ignore_viewport_layer = true;
+        }
+    }
     if (!profile.support_vertex_instance_id && (info.loads_instance_id || info.loads_vertex_id)) {
         ctx.AddExtension("SPV_KHR_shader_draw_parameters");
         ctx.AddCapability(spv::Capability::DrawParameters);
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
index 2eaeb29de..e42407f1f 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
@@ -29,7 +29,7 @@ std::optional<AttrInfo> AttrTypes(EmitContext& ctx, u32 index) {
     throw InvalidArgument("Invalid attribute type {}", type);
 }
 
-Id OutputAttrPointer(EmitContext& ctx, IR::Attribute attr) {
+std::optional<Id> OutputAttrPointer(EmitContext& ctx, IR::Attribute attr) {
     const u32 element{static_cast<u32>(attr) % 4};
     const auto element_id{[&] { return ctx.Constant(ctx.U32[1], element); }};
     if (IR::IsGeneric(attr)) {
@@ -57,6 +57,8 @@ Id OutputAttrPointer(EmitContext& ctx, IR::Attribute attr) {
         const Id clip_num{ctx.Constant(ctx.U32[1], index)};
         return ctx.OpAccessChain(ctx.output_f32, ctx.clip_distances, clip_num);
     }
+    case IR::Attribute::ViewportIndex:
+        return ctx.ignore_viewport_layer ? std::nullopt : std::optional<Id>{ctx.viewport_index};
     default:
         throw NotImplementedException("Read attribute {}", attr);
     }
@@ -204,7 +206,11 @@ Id EmitGetAttribute(EmitContext& ctx, IR::Attribute attr) {
 }
 
 void EmitSetAttribute(EmitContext& ctx, IR::Attribute attr, Id value) {
-    ctx.OpStore(OutputAttrPointer(ctx, attr), value);
+    auto output = OutputAttrPointer(ctx, attr);
+    if (!output) {
+        return;
+    }
+    ctx.OpStore(*output, value);
 }
 
 void EmitGetAttributeIndexed(EmitContext&) {
diff --git a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
index 50ffc4c19..514de6838 100644
--- a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
+++ b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
@@ -81,6 +81,9 @@ void SetAttribute(Info& info, IR::Attribute attribute) {
     case IR::Attribute::ClipDistance7:
         info.stores_clip_distance = true;
         break;
+    case IR::Attribute::ViewportIndex:
+        info.stores_viewport_index = true;
+        break;
     default:
         throw NotImplementedException("Set attribute {}", attribute);
     }
diff --git a/src/shader_recompiler/profile.h b/src/shader_recompiler/profile.h
index f4b94896c..e13cb948a 100644
--- a/src/shader_recompiler/profile.h
+++ b/src/shader_recompiler/profile.h
@@ -35,6 +35,7 @@ struct Profile {
     bool support_fp64_signed_zero_nan_preserve{};
     bool support_explicit_workgroup_layout{};
     bool support_vote{};
+    bool support_viewport_index_layer_non_geometry{};
     bool warp_size_potentially_larger_than_guest{};
 
     // FClamp is broken and OpFMax + OpFMin should be used instead
diff --git a/src/shader_recompiler/shader_info.h b/src/shader_recompiler/shader_info.h
index a62ad1e79..a5bff40bb 100644
--- a/src/shader_recompiler/shader_info.h
+++ b/src/shader_recompiler/shader_info.h
@@ -82,6 +82,7 @@ struct Info {
     bool stores_position{};
     bool stores_point_size{};
     bool stores_clip_distance{};
+    bool stores_viewport_index{};
 
     bool uses_fp16{};
     bool uses_fp64{};