[mlir][dxsa] Add conditional instructions

mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td

@@ -255,6 +255,22 @@ def DXSA_SystemValueNameAttr :
   let assemblyFormat = "`<` $value `>`";
 }
 
+def DXSA_Condition_Zero    : I32EnumAttrCase<"zero",     0>;
+def DXSA_Condition_NonZero : I32EnumAttrCase<"non_zero", 1>;
+
+def DXSA_Condition : I32EnumAttr<
+    "Condition", "DXBC boolean condition", [
+      DXSA_Condition_Zero,
+      DXSA_Condition_NonZero
+    ]> {
+  let cppNamespace = "::mlir::dxsa";
+  let genSpecializedAttr = 0;
+}
+
+def DXSA_ConditionAttr : EnumAttr<DXSADialect, DXSA_Condition, "condition"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // DXSA ComponentMask bit-enum (mask field of operand, normalized to bits 0..3)
 //===----------------------------------------------------------------------===//
@@ -1072,4 +1088,93 @@ def DXSA_DclStream : DXSA_Op<"dcl_stream"> {
   let assemblyFormat = [{ $index attr-dict }];
 }
 
+// Base class for operations with a condition attribute and a single operand.
+class DXSA_CondOp<string mnemonic> : DXSA_Op<mnemonic> {
+  let arguments = (ins DXSA_ConditionAttr:$cond, DXSA_OperandType:$operand);
+  let assemblyFormat = [{ $cond `,` $operand attr-dict}];
+}
+
+def DXSA_If : DXSA_CondOp<"if"> {
+  let summary = "branch based on logical OR result";
+  let description = [{
+    Example:
+    ```mlir
+    module {
+      %0 = dxsa.index.imm {imm = 0 : i32}
+      %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+      dxsa.if <non_zero>, %1
+    }
+    ```
+  }];
+}
+
+def DXSA_Breakc : DXSA_CondOp<"breakc"> {
+  let summary = "break (conditional)";
+  let description = [{
+    Conditionally move the point of execution to the instruction after the next
+    endloop or endswitch.
+
+    Example:
+    ```mlir
+    module {
+      %0 = dxsa.index.imm {imm = 0 : i32}
+      %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+      dxsa.breakc <non_zero>, %1
+    }
+    ```
+  }];
+}
+
+def DXSA_Continuec : DXSA_CondOp<"continuec"> {
+  let summary = "continue (conditional)";
+  let description = [{
+    Conditionally continue execution at the beginning of the current loop.
+
+    Example:
+    ```mlir
+    module {
+      %0 = dxsa.index.imm {imm = 0 : i32}
+      %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+      dxsa.continuec <non_zero>, %1
+    }
+    ```
+  }];
+}
+
+def DXSA_Retc : DXSA_CondOp<"retc"> {
+  let summary = "return (conditional)";
+  let description = [{
+    If within a subroutine, conditionally return to the instruction after the
+    call. If not inside a subroutine, terminate program execution.
+
+    Example:
+    ```mlir
+    module {
+      %0 = dxsa.index.imm {imm = 0 : i32}
+      %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+      dxsa.retc <non_zero>, %1
+    }
+    ```
+  }];
+}
+
+def DXSA_Discard : DXSA_CondOp<"discard"> {
+  let summary = "conditionally flag results of Pixel Shader to be discarded when the end of the program is reached";
+  let description = [{
+    The `discard` instruction flags the current pixel as terminated, while
+    continuing execution, so that other pixels executing in parallel may obtain
+    derivatives if necessary. Even though execution continues, all Pixel Shader
+    output writes before or after the `discard` instruction are discarded.
+
+    Example:
+    ```mlir
+    module {
+      %0 = dxsa.index.imm {imm = 0 : i32}
+      %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+      dxsa.discard <non_zero>, %1
+    }
+    ```
+  }];
+}
+
 #endif // DXSA_OPS

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -733,6 +733,12 @@ class DXBuilder {
         optionalToAttr(space));
   }
 
+  template <typename Inst>
+  Instruction buildConditionInstruction(dxsa::Condition cond, Operand operand,
+                                        Location loc) {
+    return Inst::create(builder, loc, cond, operand);
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -1476,6 +1482,19 @@ class Parser {
     return builder.buildDclSampler(id, lbound, ubound, space, *mode, loc);
   }
 
+  template <typename Inst>
+  FailureOr<Instruction> parseConditionInstruction(uint32_t opcodeToken,
+                                                   Location loc) {
+    dxsa::Condition cond = DECODE_D3D10_SB_INSTRUCTION_TEST_BOOLEAN(opcodeToken)
+                               ? dxsa::Condition::non_zero
+                               : dxsa::Condition::zero;
+
+    auto operand = parseOperand();
+    FAILURE_IF_FAILED(operand);
+
+    return builder.buildConditionInstruction<Inst>(cond, *operand, loc);
+  }
+
   OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                           Instruction &out) {
     FailureOr<Instruction> result;
@@ -1567,6 +1586,21 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_SAMPLER:
       result = parseDclSampler(opcodeToken, loc);
       break;
+    case D3D10_SB_OPCODE_IF:
+      result = parseConditionInstruction<dxsa::If>(opcodeToken, loc);
+      break;
+    case D3D10_SB_OPCODE_BREAKC:
+      result = parseConditionInstruction<dxsa::Breakc>(opcodeToken, loc);
+      break;
+    case D3D10_SB_OPCODE_CONTINUEC:
+      result = parseConditionInstruction<dxsa::Continuec>(opcodeToken, loc);
+      break;
+    case D3D10_SB_OPCODE_RETC:
+      result = parseConditionInstruction<dxsa::Retc>(opcodeToken, loc);
+      break;
+    case D3D10_SB_OPCODE_DISCARD:
+      result = parseConditionInstruction<dxsa::Discard>(opcodeToken, loc);
+      break;
     default:
       return std::nullopt;
     }

mlir/test/Target/DXSA/breakc.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/breakc.bin | FileCheck %s
+
+// CHECK: module {
+// CHECK-NEXT:   %0 = dxsa.index.imm {imm = 0 : i32}
+// CHECK-NEXT:   %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+// CHECK-NEXT:   dxsa.breakc <non_zero>, %1
+// CHECK-NEXT: }

mlir/test/Target/DXSA/continuec.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/continuec.bin | FileCheck %s
+
+// CHECK: module {
+// CHECK-NEXT:   %0 = dxsa.index.imm {imm = 0 : i32}
+// CHECK-NEXT:   %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+// CHECK-NEXT:   dxsa.continuec <non_zero>, %1
+// CHECK-NEXT: }

mlir/test/Target/DXSA/discard.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/discard.bin | FileCheck %s
+
+// CHECK: module {
+// CHECK-NEXT:   %0 = dxsa.index.imm {imm = 0 : i32}
+// CHECK-NEXT:   %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+// CHECK-NEXT:   dxsa.discard <non_zero>, %1
+// CHECK-NEXT: }

mlir/test/Target/DXSA/if.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/if.bin | FileCheck %s
+
+// CHECK: module {
+// CHECK-NEXT:   %0 = dxsa.index.imm {imm = 0 : i32}
+// CHECK-NEXT:   %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+// CHECK-NEXT:   dxsa.if <non_zero>, %1
+// CHECK-NEXT: }

mlir/test/Target/DXSA/retc.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/retc.bin | FileCheck %s
+
+// CHECK: module {
+// CHECK-NEXT:   %0 = dxsa.index.imm {imm = 0 : i32}
+// CHECK-NEXT:   %1 = dxsa.operand %0 {num_components = 4 : i32, one = 0 : i32, type = 0 : i32}
+// CHECK-NEXT:   dxsa.retc <non_zero>, %1
+// CHECK-NEXT: }