diff --git a/source/opt/amd_ext_to_khr.cpp b/source/opt/amd_ext_to_khr.cpp
index 0d84ef3..1cb5ba5 100644
--- a/source/opt/amd_ext_to_khr.cpp
+++ b/source/opt/amd_ext_to_khr.cpp
@@ -79,16 +79,121 @@
   };
 }
 
-FoldingRule NotImplementedYet() {
-  return [](IRContext*, Instruction*,
+// Returns a folding rule that will replace the SwizzleInvocationsAMD extended
+// instruction in the SPV_AMD_shader_ballot extension.
+//
+// The instruction
+//
+//  %offset = OpConstantComposite %v3uint %x %y %z %w
+//  %result = OpExtInst %type %1 SwizzleInvocationsAMD %data %offset
+//
+// is replaced with
+//
+// potentially new constants and types
+//
+// clang-format off
+//         %uint_max = OpConstant %uint 0xFFFFFFFF
+//           %v4uint = OpTypeVector %uint 4
+//     %ballot_value = OpConstantComposite %v4uint %uint_max %uint_max %uint_max %uint_max
+//             %null = OpConstantNull %type
+// clang-format on
+//
+// and the following code in the function body
+//
+// clang-format off
+//         %id = OpLoad %uint %SubgroupLocalInvocationId
+//   %quad_idx = OpBitwiseAnd %uint %id %uint_3
+//   %quad_ldr = OpBitwiseXor %uint %id %quad_idx
+//  %my_offset = OpVectorExtractDynamic %uint %offset %quad_idx
+// %target_inv = OpIAdd %uint %quad_ldr %my_offset
+//  %is_active = OpGroupNonUniformBallotBitExtract %bool %uint_3 %ballot_value %target_inv
+//    %shuffle = OpGroupNonUniformShuffle %type %uint_3 %data %target_inv
+//     %result = OpSelect %type %is_active %shuffle %null
+// clang-format on
+//
+// Also adding the capabilities and builtins that are needed.
+FoldingRule ReplaceSwizzleInvocations() {
+  return [](IRContext* ctx, Instruction* inst,
             const std::vector<const analysis::Constant*>&) {
-    assert(false && "Replacement not implemented yet.");
-    return false;
+    analysis::TypeManager* type_mgr = ctx->get_type_mgr();
+    analysis::ConstantManager* const_mgr = ctx->get_constant_mgr();
+
+    ctx->AddExtension("SPV_KHR_shader_ballot");
+    ctx->AddCapability(SpvCapabilityGroupNonUniformBallot);
+    ctx->AddCapability(SpvCapabilityGroupNonUniformShuffle);
+
+    InstructionBuilder ir_builder(
+        ctx, inst,
+        IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
+
+    uint32_t data_id = inst->GetSingleWordInOperand(2);
+    uint32_t offset_id = inst->GetSingleWordInOperand(3);
+
+    // Get the subgroup invocation id.
+    uint32_t var_id =
+        ctx->GetBuiltinInputVarId(SpvBuiltInSubgroupLocalInvocationId);
+    assert(var_id != 0 && "Could not get SubgroupLocalInvocationId variable.");
+    Instruction* var_inst = ctx->get_def_use_mgr()->GetDef(var_id);
+    Instruction* var_ptr_type =
+        ctx->get_def_use_mgr()->GetDef(var_inst->type_id());
+    uint32_t uint_type_id = var_ptr_type->GetSingleWordInOperand(1);
+
+    Instruction* id = ir_builder.AddLoad(uint_type_id, var_id);
+
+    uint32_t quad_mask = ir_builder.GetUintConstantId(3);
+
+    // This gives the offset in the group of 4 of this invocation.
+    Instruction* quad_idx = ir_builder.AddBinaryOp(
+        uint_type_id, SpvOpBitwiseAnd, id->result_id(), quad_mask);
+
+    // Get the invocation id of the first invocation in the group of 4.
+    Instruction* quad_ldr = ir_builder.AddBinaryOp(
+        uint_type_id, SpvOpBitwiseXor, id->result_id(), quad_idx->result_id());
+
+    // Get the offset of the target invocation from the offset vector.
+    Instruction* my_offset =
+        ir_builder.AddBinaryOp(uint_type_id, SpvOpVectorExtractDynamic,
+                               offset_id, quad_idx->result_id());
+
+    // Determine the index of the invocation to read from.
+    Instruction* target_inv = ir_builder.AddBinaryOp(
+        uint_type_id, SpvOpIAdd, quad_ldr->result_id(), my_offset->result_id());
+
+    // Do the group operations
+    uint32_t uint_max_id = ir_builder.GetUintConstantId(0xFFFFFFFF);
+    uint32_t subgroup_scope = ir_builder.GetUintConstantId(SpvScopeSubgroup);
+    const auto* ballot_value_const = const_mgr->GetConstant(
+        type_mgr->GetUIntVectorType(4),
+        {uint_max_id, uint_max_id, uint_max_id, uint_max_id});
+    Instruction* ballot_value =
+        const_mgr->GetDefiningInstruction(ballot_value_const);
+    Instruction* is_active = ir_builder.AddNaryOp(
+        type_mgr->GetBoolTypeId(), SpvOpGroupNonUniformBallotBitExtract,
+        {subgroup_scope, ballot_value->result_id(), target_inv->result_id()});
+    Instruction* shuffle = ir_builder.AddNaryOp(
+        inst->type_id(), SpvOpGroupNonUniformShuffle,
+        {subgroup_scope, data_id, target_inv->result_id()});
+
+    // Create the null constant to use in the select.
+    const auto* null = const_mgr->GetConstant(
+        type_mgr->GetType(inst->type_id()), std::vector<uint32_t>());
+    Instruction* null_inst = const_mgr->GetDefiningInstruction(null);
+
+    // Build the select.
+    inst->SetOpcode(SpvOpSelect);
+    Instruction::OperandList new_operands;
+    new_operands.push_back({SPV_OPERAND_TYPE_ID, {is_active->result_id()}});
+    new_operands.push_back({SPV_OPERAND_TYPE_ID, {shuffle->result_id()}});
+    new_operands.push_back({SPV_OPERAND_TYPE_ID, {null_inst->result_id()}});
+
+    inst->SetInOperands(std::move(new_operands));
+    ctx->UpdateDefUse(inst);
+    return true;
   };
 }
 
-// Returns a folding rule that will replace the WriteInvocationAMD extended
-// instruction in the SPV_AMD_shader_ballot extension.
+// Returns a folding rule that will replace the SwizzleInvocationsMaskedAMD
+// extended instruction in the SPV_AMD_shader_ballot extension.
 //
 // The instruction
 //
@@ -120,7 +225,7 @@
 // clang-format on
 //
 // Also adding the capabilities and builtins that are needed.
-FoldingRule FoldSwizzleInvocationsMasked() {
+FoldingRule ReplaceSwizzleInvocationsMasked() {
   return [](IRContext* ctx, Instruction* inst,
             const std::vector<const analysis::Constant*>&) {
     analysis::TypeManager* type_mgr = ctx->get_type_mgr();
@@ -353,9 +458,9 @@
         context()->module()->GetExtInstImportId("SPV_AMD_shader_ballot");
 
     ext_rules_[{extension_id, AmdShaderBallotSwizzleInvocationsAMD}].push_back(
-        NotImplementedYet());
+        ReplaceSwizzleInvocations());
     ext_rules_[{extension_id, AmdShaderBallotSwizzleInvocationsMaskedAMD}]
-        .push_back(FoldSwizzleInvocationsMasked());
+        .push_back(ReplaceSwizzleInvocationsMasked());
     ext_rules_[{extension_id, AmdShaderBallotWriteInvocationAMD}].push_back(
         ReplaceWriteInvocation());
     ext_rules_[{extension_id, AmdShaderBallotMbcntAMD}].push_back(
diff --git a/test/opt/amd_ext_to_khr.cpp b/test/opt/amd_ext_to_khr.cpp
index b260761..7a6d4b4 100644
--- a/test/opt/amd_ext_to_khr.cpp
+++ b/test/opt/amd_ext_to_khr.cpp
@@ -135,6 +135,54 @@
   SinglePassRunAndMatch<AmdExtensionToKhrPass>(text, true);
 }
 
+TEST_F(AmdExtToKhrTest, ReplaceSwizzleInvocationsAMD) {
+  const std::string text = R"(
+; CHECK: OpCapability Shader
+; CHECK-NOT: OpExtension "SPV_AMD_shader_ballot"
+; CHECK-NOT: OpExtInstImport "SPV_AMD_shader_ballot"
+; CHECK: OpDecorate [[var:%\w+]] BuiltIn SubgroupLocalInvocationId
+; CHECK: [[subgroup:%\w+]] = OpConstant %uint 3
+; CHECK: [[offset:%\w+]] = OpConstantComposite %v4uint
+; CHECK: [[var]] = OpVariable %_ptr_Input_uint Input
+; CHECK: [[uint_max:%\w+]] = OpConstant %uint 4294967295
+; CHECK: [[ballot_value:%\w+]] = OpConstantComposite %v4uint [[uint_max]] [[uint_max]] [[uint_max]] [[uint_max]]
+; CHECK: [[null:%\w+]] = OpConstantNull [[type:%\w+]]
+; CHECK: OpFunction
+; CHECK-NEXT: OpLabel
+; CHECK-NEXT: [[data:%\w+]] = OpUndef [[type]]
+; CHECK-NEXT: [[id:%\w+]] = OpLoad %uint [[var]]
+; CHECK-NEXT: [[quad_idx:%\w+]] = OpBitwiseAnd %uint [[id]] %uint_3
+; CHECK-NEXT: [[quad_ldr:%\w+]] = OpBitwiseXor %uint [[id]] [[quad_idx]]
+; CHECK-NEXT: [[my_offset:%\w+]] = OpVectorExtractDynamic %uint [[offset]] [[quad_idx]]
+; CHECK-NEXT: [[target_inv:%\w+]] = OpIAdd %uint [[quad_ldr]] [[my_offset]]
+; CHECK-NEXT: [[is_active:%\w+]] = OpGroupNonUniformBallotBitExtract %bool [[subgroup]] [[ballot_value]] [[target_inv]]
+; CHECK-NEXT: [[shuffle:%\w+]] = OpGroupNonUniformShuffle [[type]] [[subgroup]] [[data]] [[target_inv]]
+; CHECK-NEXT: [[result:%\w+]] = OpSelect [[type]] [[is_active]] [[shuffle]] [[null]]
+               OpCapability Shader
+               OpExtension "SPV_AMD_shader_ballot"
+        %ext = OpExtInstImport "SPV_AMD_shader_ballot"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %1 "func"
+               OpExecutionMode %1 OriginUpperLeft
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_x = OpConstant %uint 1
+     %uint_y = OpConstant %uint 2
+     %uint_z = OpConstant %uint 3
+     %uint_w = OpConstant %uint 0
+     %v4uint = OpTypeVector %uint 4
+     %offset = OpConstantComposite %v4uint %uint_x %uint_y %uint_z %uint_x
+          %1 = OpFunction %void None %3
+          %6 = OpLabel
+          %data = OpUndef %uint
+          %9 = OpExtInst %uint %ext SwizzleInvocationsAMD %data %offset
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<AmdExtensionToKhrPass>(text, true);
+}
 TEST_F(AmdExtToKhrTest, ReplaceSwizzleInvocationsMaskedAMD) {
   const std::string text = R"(
 ; CHECK: OpCapability Shader