From 26966287c797a9caf858dec4fa4d491a56caa6be Mon Sep 17 00:00:00 2001
From: PixelyIon <pixelyion@protonmail.com>
Date: Fri, 12 Nov 2021 00:36:05 +0530
Subject: [PATCH] Implement Maxwell3D Shader Program Registers

This commit added basic shader program registers, they simply track the address a shader is pointed to at the moment. No parsing of the shader program is done within them.
---
 .../gpu/interconnect/graphics_context.h       | 36 +++++++++++++++++++
 .../skyline/soc/gm20b/engines/maxwell/types.h | 32 +++++++++++++++++
 .../skyline/soc/gm20b/engines/maxwell_3d.cpp  | 16 +++++++++
 .../skyline/soc/gm20b/engines/maxwell_3d.h    |  3 ++
 4 files changed, 87 insertions(+)

diff --git a/app/src/main/cpp/skyline/gpu/interconnect/graphics_context.h b/app/src/main/cpp/skyline/gpu/interconnect/graphics_context.h
index 93fb5d96..c772f515 100644
--- a/app/src/main/cpp/skyline/gpu/interconnect/graphics_context.h
+++ b/app/src/main/cpp/skyline/gpu/interconnect/graphics_context.h
@@ -326,5 +326,41 @@ namespace skyline::gpu::interconnect {
             scissor.offset.y = bounds.minimum;
             scissor.extent.height = bounds.maximum - bounds.minimum;
         }
+
+        /* Shader Program */
+      private:
+        struct Shader {
+            bool enabled{false};
+            u32 offset{}; //!< Offset of the shader from the base IOVA
+            span<u8> data; //!< The shader bytecode in the CPU AS
+        };
+
+        u64 shaderBaseIova{}; //!< The base IOVA that shaders are located at an offset from
+        std::array<Shader, maxwell3d::StageCount> boundShaders{};
+
+      public:
+        void SetShaderBaseIovaHigh(u32 high) {
+            shaderBaseIova = (shaderBaseIova & std::numeric_limits<u32>::max()) | (static_cast<u64>(high) << 32);
+            for (auto &shader : boundShaders)
+                shader.data = span<u8>{};
+        }
+
+        void SetShaderBaseIovaLow(u32 low) {
+            shaderBaseIova = (shaderBaseIova & (static_cast<u64>(std::numeric_limits<u32>::max()) << 32)) | low;
+            for (auto &shader : boundShaders)
+                shader.data = span<u8>{};
+        }
+
+        void SetShaderEnabled(maxwell3d::StageId stage, bool enabled) {
+            auto &shader{boundShaders[static_cast<size_t>(stage)]};
+            shader.enabled = enabled;
+            shader.data = span<u8>{};
+        }
+
+        void SetShaderOffset(maxwell3d::StageId stage, u32 offset) {
+            auto &shader{boundShaders[static_cast<size_t>(stage)]};
+            shader.offset = offset;
+            shader.data = span<u8>{};
+        }
     };
 }
diff --git a/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell/types.h b/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell/types.h
index c2e3cfbf..060c0dd5 100644
--- a/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell/types.h
+++ b/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell/types.h
@@ -525,5 +525,37 @@ namespace skyline::soc::gm20b::engine::maxwell3d::type {
     };
     static_assert(sizeof(SemaphoreInfo) == sizeof(u32));
 
+    constexpr static size_t StageCount{6}; //!< Amount of pipeline stages on Maxwell 3D
+
+    /**
+     * @brief All the pipeline stages that Maxwell3D supports for draws
+     */
+    enum class StageId {
+        VertexA = 0,
+        VertexB = 1,
+        TessellationControl = 2,
+        TessellationEvaluation = 3,
+        Geometry = 4,
+        Fragment = 5,
+    };
+    static_assert(static_cast<size_t>(StageId::Fragment) + 1 == StageCount);
+
+    /**
+     * @brief The arguments to set a shader program for a pipeline stage
+     */
+    struct SetProgramInfo {
+        struct {
+            bool enable : 1;
+            u8 _pad0_ : 3;
+            StageId stage : 4;
+            u32 _pad1_ : 24;
+        } info;
+        u32 offset; //!< Offset from the base shader memory IOVA
+        u32 _pad2_;
+        u32 gprCount; //!< Amount of GPRs used by the shader program
+        u32 _pad3_[12];
+    };
+    static_assert(sizeof(SetProgramInfo) == (sizeof(u32) * 0x10));
+
     #pragma pack(pop)
 }
diff --git a/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.cpp b/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.cpp
index a35c42cb..e6e51266 100644
--- a/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.cpp
+++ b/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.cpp
@@ -226,6 +226,14 @@ namespace skyline::soc::gm20b::engine::maxwell3d {
                 MAXWELL3D_CASE(renderTargetControl, {
                     context.UpdateRenderTargetControl(renderTargetControl);
                 })
+                #define SET_SHADER_ENABLE_CALLBACK(z, index, data)     \
+                MAXWELL3D_ARRAY_STRUCT_CASE(setProgram, index, info, { \
+                    context.SetShaderEnabled(info.stage, info.enable); \
+                })
+
+                BOOST_PP_REPEAT(6, SET_SHADER_ENABLE_CALLBACK, 0)
+                static_assert(type::StageCount == 6 && type::StageCount < BOOST_PP_LIMIT_REPEAT);
+                #undef SET_SHADER_ENABLE_CALLBACK
             }
         }
 
@@ -282,6 +290,14 @@ namespace skyline::soc::gm20b::engine::maxwell3d {
                 }
             })
 
+            #define SET_SHADER_OFFSET_CALLBACK(z, index, data)                               \
+                MAXWELL3D_ARRAY_STRUCT_CASE(setProgram, index, offset, {                     \
+                    context.SetShaderOffset(registers.setProgram[index].info.stage, offset); \
+                })
+
+            BOOST_PP_REPEAT(6, SET_SHADER_OFFSET_CALLBACK, 0)
+            static_assert(type::StageCount == 6 && type::StageCount < BOOST_PP_LIMIT_REPEAT);
+
             MAXWELL3D_ARRAY_CASE(firmwareCall, 4, {
                 registers.raw[0xD00] = 1;
             })
diff --git a/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.h b/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.h
index 7498c4ab..3872a23f 100644
--- a/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.h
+++ b/app/src/main/cpp/skyline/soc/gm20b/engines/maxwell_3d.h
@@ -200,6 +200,9 @@ namespace skyline::soc::gm20b::engine::maxwell3d {
             Register<0x6C0, Semaphore> semaphore;
 
             Register<0x780, std::array<type::Blend, type::RenderTargetCount>> independentBlend;
+
+            Register<0x800, std::array<type::SetProgramInfo, type::StageCount>> setProgram;
+
             Register<0x8C0, u32[0x20]> firmwareCall;
         };
         static_assert(sizeof(Registers) == (RegisterCount * sizeof(u32)));