11 files changed, 178 insertions, 43 deletions

diff --git a/src/shader_recompiler/frontend/maxwell/control_flow.cpp b/src/shader_recompiler/frontend/maxwell/control_flow.cpp
index 715c0e92d..4f6707fae 100644
--- a/src/shader_recompiler/frontend/maxwell/control_flow.cpp
+++ b/src/shader_recompiler/frontend/maxwell/control_flow.cpp
@@ -104,6 +104,7 @@ bool HasFlowTest(Opcode opcode) {
     case Opcode::EXIT:
     case Opcode::JMP:
     case Opcode::JMX:
+    case Opcode::KIL:
     case Opcode::BRK:
     case Opcode::CONT:
     case Opcode::LONGJMP:
@@ -287,6 +288,13 @@ CFG::AnalysisState CFG::AnalyzeInst(Block* block, FunctionId function_id, Locati
         block->end = pc;
         return AnalysisState::Branch;
     }
+    case Opcode::KIL: {
+        const Predicate pred{inst.Pred()};
+        const auto ir_pred{static_cast<IR::Pred>(pred.index)};
+        const IR::Condition cond{inst.branch.flow_test, ir_pred, pred.negated};
+        AnalyzeCondInst(block, function_id, pc, EndClass::Kill, cond);
+        return AnalysisState::Branch;
+    }
     case Opcode::PBK:
     case Opcode::PCNT:
     case Opcode::PEXIT:
@@ -324,13 +332,12 @@ CFG::AnalysisState CFG::AnalyzeInst(Block* block, FunctionId function_id, Locati
         return AnalysisState::Continue;
     }
     const IR::Condition cond{static_cast<IR::Pred>(pred.index), pred.negated};
-    AnalyzeCondInst(block, function_id, pc, EndClass::Branch, cond, true);
+    AnalyzeCondInst(block, function_id, pc, EndClass::Branch, cond);
     return AnalysisState::Branch;
 }
 
 void CFG::AnalyzeCondInst(Block* block, FunctionId function_id, Location pc,
-                          EndClass insn_end_class, IR::Condition cond,
-                          bool visit_conditional_inst) {
+                          EndClass insn_end_class, IR::Condition cond) {
     if (block->begin != pc) {
         // If the block doesn't start in the conditional instruction
         // mark it as a label to visit it later
@@ -356,14 +363,16 @@ void CFG::AnalyzeCondInst(Block* block, FunctionId function_id, Location pc,
     // Impersonate the visited block with a virtual block
     *block = std::move(virtual_block);
     // Set the end properties of the conditional instruction
-    conditional_block->end = visit_conditional_inst ? (pc + 1) : pc;
+    conditional_block->end = pc + 1;
     conditional_block->end_class = insn_end_class;
     // Add a label to the instruction after the conditional instruction
     Block* const endif_block{AddLabel(conditional_block, block->stack, pc + 1, function_id)};
     // Branch to the next instruction from the virtual block
     block->branch_false = endif_block;
-    // And branch to it from the conditional instruction if it is a branch
-    if (insn_end_class == EndClass::Branch) {
+    // And branch to it from the conditional instruction if it is a branch or a kill instruction
+    // Kill instructions are considered a branch because they demote to a helper invocation and
+    // execution may continue.
+    if (insn_end_class == EndClass::Branch || insn_end_class == EndClass::Kill) {
         conditional_block->cond = IR::Condition{true};
         conditional_block->branch_true = endif_block;
         conditional_block->branch_false = nullptr;
@@ -415,7 +424,7 @@ CFG::AnalysisState CFG::AnalyzeEXIT(Block* block, FunctionId function_id, Locati
             throw NotImplementedException("Conditional EXIT with PEXIT token");
         }
         const IR::Condition cond{flow_test, static_cast<IR::Pred>(pred.index), pred.negated};
-        AnalyzeCondInst(block, function_id, pc, EndClass::Exit, cond, false);
+        AnalyzeCondInst(block, function_id, pc, EndClass::Exit, cond);
         return AnalysisState::Branch;
     }
     if (const std::optional<Location> exit_pc{block->stack.Peek(Token::PEXIT)}) {
@@ -425,7 +434,7 @@ CFG::AnalysisState CFG::AnalyzeEXIT(Block* block, FunctionId function_id, Locati
         block->branch_false = nullptr;
         return AnalysisState::Branch;
     }
-    block->end = pc;
+    block->end = pc + 1;
     block->end_class = EndClass::Exit;
     return AnalysisState::Branch;
 }
@@ -505,6 +514,12 @@ std::string CFG::Dot() const {
                                    node_uid);
                 ++node_uid;
                 break;
+            case EndClass::Kill:
+                dot += fmt::format("\t\t{}->N{};\n", name, node_uid);
+                dot += fmt::format("\t\tN{} [label=\"Kill\"][shape=square][style=stripped];\n",
+                                   node_uid);
+                ++node_uid;
+                break;
             }
         }
         if (function.entrypoint == 8) {
diff --git a/src/shader_recompiler/frontend/maxwell/control_flow.h b/src/shader_recompiler/frontend/maxwell/control_flow.h
index fe74f210f..22f134194 100644
--- a/src/shader_recompiler/frontend/maxwell/control_flow.h
+++ b/src/shader_recompiler/frontend/maxwell/control_flow.h
@@ -29,6 +29,7 @@ enum class EndClass {
     Call,
     Exit,
     Return,
+    Kill,
 };
 
 enum class Token {
@@ -130,7 +131,7 @@ private:
     AnalysisState AnalyzeInst(Block* block, FunctionId function_id, Location pc);
 
     void AnalyzeCondInst(Block* block, FunctionId function_id, Location pc, EndClass insn_end_class,
-                         IR::Condition cond, bool visit_conditional_inst);
+                         IR::Condition cond);
 
     /// Return true when the branch instruction is confirmed to be a branch
     bool AnalyzeBranch(Block* block, FunctionId function_id, Location pc, Instruction inst,
diff --git a/src/shader_recompiler/frontend/maxwell/program.cpp b/src/shader_recompiler/frontend/maxwell/program.cpp
index 8bfa64326..0074eb89b 100644
--- a/src/shader_recompiler/frontend/maxwell/program.cpp
+++ b/src/shader_recompiler/frontend/maxwell/program.cpp
@@ -32,6 +32,7 @@ IR::Program TranslateProgram(ObjectPool<IR::Inst>& inst_pool, ObjectPool<IR::Blo
     IR::Program program;
     program.blocks = VisitAST(inst_pool, block_pool, env, cfg);
     program.post_order_blocks = PostOrder(program.blocks);
+    program.stage = env.ShaderStage();
     RemoveUnreachableBlocks(program);
 
     // Replace instructions before the SSA rewrite
diff --git a/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp b/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp
index 5f5d9cf17..cec03e73e 100644
--- a/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp
+++ b/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp
@@ -45,6 +45,7 @@ enum class StatementType {
     Loop,
     Break,
     Return,
+    Kill,
     Function,
     Identity,
     Not,
@@ -70,6 +71,7 @@ struct If {};
 struct Loop {};
 struct Break {};
 struct Return {};
+struct Kill {};
 struct FunctionTag {};
 struct Identity {};
 struct Not {};
@@ -93,6 +95,7 @@ struct Statement : ListBaseHook {
     Statement(Break, Statement* cond_, Statement* up_)
         : cond{cond_}, up{up_}, type{StatementType::Break} {}
     Statement(Return) : type{StatementType::Return} {}
+    Statement(Kill) : type{StatementType::Kill} {}
     Statement(FunctionTag) : children{}, type{StatementType::Function} {}
     Statement(Identity, IR::Condition cond_) : guest_cond{cond_}, type{StatementType::Identity} {}
     Statement(Not, Statement* op_) : op{op_}, type{StatementType::Not} {}
@@ -174,6 +177,9 @@ std::string DumpTree(const Tree& tree, u32 indentation = 0) {
         case StatementType::Return:
             ret += fmt::format("{}    return;\n", indent);
             break;
+        case StatementType::Kill:
+            ret += fmt::format("{}    kill;\n", indent);
+            break;
         case StatementType::SetVariable:
             ret += fmt::format("{}    goto_L{} = {};\n", indent, stmt->id, DumpExpr(stmt->op));
             break;
@@ -424,6 +430,9 @@ private:
                 gotos.push_back(root.insert(ip, *goto_stmt));
                 break;
             }
+            case Flow::EndClass::Kill:
+                root.insert(ip, *pool.Create(Kill{}));
+                break;
             }
         }
     }
@@ -729,6 +738,15 @@ private:
                 current_block = nullptr;
                 break;
             }
+            case StatementType::Kill: {
+                if (!current_block) {
+                    current_block = block_pool.Create(inst_pool);
+                    block_list.push_back(current_block);
+                }
+                IR::IREmitter{*current_block}.DemoteToHelperInvocation(continue_block);
+                current_block = nullptr;
+                break;
+            }
             default:
                 throw NotImplementedException("Statement type {}", stmt.type);
             }
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/exit.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/exit.cpp
deleted file mode 100644
index e98bbd0d1..000000000
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/exit.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-// Copyright 2021 yuzu Emulator Project
-// Licensed under GPLv2 or any later version
-// Refer to the license.txt file included.
-
-#include "common/common_types.h"
-#include "shader_recompiler/exception.h"
-#include "shader_recompiler/frontend/maxwell/translate/impl/impl.h"
-
-namespace Shader::Maxwell {
-
-void TranslatorVisitor::EXIT(u64) {
-    ir.Exit();
-}
-
-} // namespace Shader::Maxwell
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/exit_program.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/exit_program.cpp
new file mode 100644
index 000000000..ea9b33da9
--- /dev/null
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/exit_program.cpp
@@ -0,0 +1,43 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include "common/common_types.h"
+#include "shader_recompiler/exception.h"
+#include "shader_recompiler/frontend/maxwell/translate/impl/impl.h"
+
+namespace Shader::Maxwell {
+namespace {
+void ExitFragment(TranslatorVisitor& v) {
+    const ProgramHeader sph{v.env.SPH()};
+    IR::Reg src_reg{IR::Reg::R0};
+    for (u32 render_target = 0; render_target < 8; ++render_target) {
+        const std::array<bool, 4> mask{sph.ps.EnabledOutputComponents(render_target)};
+        for (u32 component = 0; component < 4; ++component) {
+            if (!mask[component]) {
+                continue;
+            }
+            v.ir.SetFragColor(render_target, component, v.F(src_reg));
+            ++src_reg;
+        }
+    }
+    if (sph.ps.omap.sample_mask != 0) {
+        throw NotImplementedException("Sample mask");
+    }
+    if (sph.ps.omap.depth != 0) {
+        throw NotImplementedException("Fragment depth");
+    }
+}
+} // Anonymous namespace
+
+void TranslatorVisitor::EXIT() {
+    switch (env.ShaderStage()) {
+    case Stage::Fragment:
+        ExitFragment(*this);
+        break;
+    default:
+        break;
+    }
+}
+
+} // namespace Shader::Maxwell
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h b/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h
index e3e298c3b..ed81d9c36 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h
@@ -108,7 +108,7 @@ public:
     void DSETP_reg(u64 insn);
     void DSETP_cbuf(u64 insn);
     void DSETP_imm(u64 insn);
-    void EXIT(u64 insn);
+    void EXIT();
     void F2F_reg(u64 insn);
     void F2F_cbuf(u64 insn);
     void F2F_imm(u64 insn);
@@ -220,7 +220,7 @@ public:
     void JCAL(u64 insn);
     void JMP(u64 insn);
     void JMX(u64 insn);
-    void KIL(u64 insn);
+    void KIL();
     void LD(u64 insn);
     void LDC(u64 insn);
     void LDG(u64 insn);
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp
index ad97786d4..2922145ee 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp
@@ -11,6 +11,13 @@
 
 namespace Shader::Maxwell {
 namespace {
+enum class Size : u64 {
+    B32,
+    B64,
+    B96,
+    B128,
+};
+
 enum class InterpolationMode : u64 {
     Pass,
     Multiply,
@@ -23,8 +30,85 @@ enum class SampleMode : u64 {
     Centroid,
     Offset,
 };
+
+int NumElements(Size size) {
+    switch (size) {
+    case Size::B32:
+        return 1;
+    case Size::B64:
+        return 2;
+    case Size::B96:
+        return 3;
+    case Size::B128:
+        return 4;
+    }
+    throw InvalidArgument("Invalid size {}", size);
+}
 } // Anonymous namespace
 
+void TranslatorVisitor::ALD(u64 insn) {
+    union {
+        u64 raw;
+        BitField<0, 8, IR::Reg> dest_reg;
+        BitField<8, 8, IR::Reg> index_reg;
+        BitField<20, 10, u64> absolute_offset;
+        BitField<20, 11, s64> relative_offset;
+        BitField<39, 8, IR::Reg> stream_reg;
+        BitField<32, 1, u64> o;
+        BitField<31, 1, u64> patch;
+        BitField<47, 2, Size> size;
+    } const ald{insn};
+
+    if (ald.o != 0) {
+        throw NotImplementedException("O");
+    }
+    if (ald.patch != 0) {
+        throw NotImplementedException("P");
+    }
+    if (ald.index_reg != IR::Reg::RZ) {
+        throw NotImplementedException("Indexed");
+    }
+    const u64 offset{ald.absolute_offset.Value()};
+    if (offset % 4 != 0) {
+        throw NotImplementedException("Unaligned absolute offset {}", offset);
+    }
+    const int num_elements{NumElements(ald.size)};
+    for (int element = 0; element < num_elements; ++element) {
+        F(ald.dest_reg + element, ir.GetAttribute(IR::Attribute{offset / 4 + element}));
+    }
+}
+
+void TranslatorVisitor::AST(u64 insn) {
+    union {
+        u64 raw;
+        BitField<0, 8, IR::Reg> src_reg;
+        BitField<8, 8, IR::Reg> index_reg;
+        BitField<20, 10, u64> absolute_offset;
+        BitField<20, 11, s64> relative_offset;
+        BitField<31, 1, u64> patch;
+        BitField<39, 8, IR::Reg> stream_reg;
+        BitField<47, 2, Size> size;
+    } const ast{insn};
+
+    if (ast.patch != 0) {
+        throw NotImplementedException("P");
+    }
+    if (ast.stream_reg != IR::Reg::RZ) {
+        throw NotImplementedException("Stream store");
+    }
+    if (ast.index_reg != IR::Reg::RZ) {
+        throw NotImplementedException("Indexed store");
+    }
+    const u64 offset{ast.absolute_offset.Value()};
+    if (offset % 4 != 0) {
+        throw NotImplementedException("Unaligned absolute offset {}", offset);
+    }
+    const int num_elements{NumElements(ast.size)};
+    for (int element = 0; element < num_elements; ++element) {
+        ir.SetAttribute(IR::Attribute{offset / 4 + element}, F(ast.src_reg + element));
+    }
+}
+
 void TranslatorVisitor::IPA(u64 insn) {
     // IPA is the instruction used to read varyings from a fragment shader.
     // gl_FragCoord is mapped to the gl_Position attribute.
@@ -51,7 +135,7 @@ void TranslatorVisitor::IPA(u64 insn) {
     // }
     const bool is_indexed{ipa.idx != 0 && ipa.index_reg != IR::Reg::RZ};
     if (is_indexed) {
-        throw NotImplementedException("IPA.IDX");
+        throw NotImplementedException("IDX");
     }
 
     const IR::Attribute attribute{ipa.attribute};
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
index 9675cef54..59252bcc5 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
@@ -17,14 +17,6 @@ void TranslatorVisitor::AL2P(u64) {
     ThrowNotImplemented(Opcode::AL2P);
 }
 
-void TranslatorVisitor::ALD(u64) {
-    ThrowNotImplemented(Opcode::ALD);
-}
-
-void TranslatorVisitor::AST(u64) {
-    ThrowNotImplemented(Opcode::AST);
-}
-
 void TranslatorVisitor::ATOM_cas(u64) {
     ThrowNotImplemented(Opcode::ATOM_cas);
 }
@@ -153,10 +145,6 @@ void TranslatorVisitor::DSETP_imm(u64) {
     ThrowNotImplemented(Opcode::DSETP_imm);
 }
 
-void TranslatorVisitor::EXIT(u64) {
-    throw LogicError("Visting EXIT instruction");
-}
-
 void TranslatorVisitor::F2F_reg(u64) {
     ThrowNotImplemented(Opcode::F2F_reg);
 }
@@ -345,8 +333,8 @@ void TranslatorVisitor::JMX(u64) {
     ThrowNotImplemented(Opcode::JMX);
 }
 
-void TranslatorVisitor::KIL(u64) {
-    ThrowNotImplemented(Opcode::KIL);
+void TranslatorVisitor::KIL() {
+    // KIL is a no-op
 }
 
 void TranslatorVisitor::LD(u64) {
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch.cpp
index 98d9f4c64..0fbb87ec4 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch.cpp
@@ -215,7 +215,7 @@ void TranslatorVisitor::TEX(u64 insn) {
         BitField<36, 13, u64> cbuf_offset;
     } const tex{insn};
 
-    Impl(*this, insn, tex.aoffi != 0, tex.blod, tex.lc != 0, static_cast<u32>(tex.cbuf_offset));
+    Impl(*this, insn, tex.aoffi != 0, tex.blod, tex.lc != 0, static_cast<u32>(tex.cbuf_offset * 4));
 }
 
 void TranslatorVisitor::TEX_b(u64 insn) {
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch_swizzled.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch_swizzled.cpp
index ac1615b00..54f0df754 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch_swizzled.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/texture_fetch_swizzled.cpp
@@ -70,7 +70,7 @@ IR::F32 ReadArray(TranslatorVisitor& v, const IR::U32& value) {
 
 IR::Value Sample(TranslatorVisitor& v, u64 insn) {
     const Encoding texs{insn};
-    const IR::U32 handle{v.ir.Imm32(static_cast<u32>(texs.cbuf_offset))};
+    const IR::U32 handle{v.ir.Imm32(static_cast<u32>(texs.cbuf_offset * 4))};
     const IR::F32 zero{v.ir.Imm32(0.0f)};
     const IR::Reg reg_a{texs.src_reg_a};
     const IR::Reg reg_b{texs.src_reg_b};