[mlir][dxsa] Add dcl_sampler instruction

Example:
    dxsa.dcl_sampler <id = 3, mode = default>
    dxsa.dcl_sampler <id = 0, mode = comparison, lbound = 0, ubound = 3, space = 1>

Signed-off-by: Vladimir Shiryaev <tagolog@users.noreply.github.com>

Author: tagolog

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

@@ -299,6 +299,27 @@ def DXSA_ConstantBufferAccessPatternAttr :
   let assemblyFormat = "`<` $value `>`";
 }
 
+// `default` is a reserved C++ keyword, so the symbolic (C++ enum case) names
+// use PascalCase while the printable asm form stays lowercase.
+def DXSA_SamplerMode_Default    : I32EnumAttrCase<"Default",    0, "default">;
+def DXSA_SamplerMode_Comparison : I32EnumAttrCase<"Comparison", 1, "comparison">;
+def DXSA_SamplerMode_Mono       : I32EnumAttrCase<"Mono",       2, "mono">;
+
+def DXSA_SamplerMode : I32EnumAttr<
+    "SamplerMode", "sampler mode", [
+      DXSA_SamplerMode_Default,
+      DXSA_SamplerMode_Comparison,
+      DXSA_SamplerMode_Mono
+    ]> {
+  let cppNamespace = "::mlir::dxsa";
+  let genSpecializedAttr = 0;
+}
+
+def DXSA_SamplerModeAttr :
+    EnumAttr<DXSADialect, DXSA_SamplerMode, "sampler_mode"> {
+  let assemblyFormat = "$value";
+}
+
 //===----------------------------------------------------------------------===//
 // DXSA attribute Constraints
 //===----------------------------------------------------------------------===//
@@ -969,4 +990,32 @@ def DXSA_DclConstantBuffer : DXSA_Op<"dcl_constant_buffer"> {
   let hasVerifier = 1;
 }
 
+def DXSA_DclSampler : DXSA_Op<"dcl_sampler"> {
+  let summary = "declares a sampler that will be referenced in the shader";
+  let description = [{
+    The `dxsa.dcl_sampler` operation declares sampler that will be referenced in the shader.
+
+    Examples:
+
+    ```mlir
+    dxsa.dcl_sampler <id = 3, mode = default>
+    dxsa.dcl_sampler <id = 0, mode = comparison, lbound = 0, ubound = 3, space = 1>
+    ```
+  }];
+
+  let arguments = (ins
+      I32Attr:$id,
+      DXSA_SamplerModeAttr:$mode,
+      OptionalAttr<I32Attr>:$lbound,
+      OptionalAttr<I32Attr>:$ubound,
+      OptionalAttr<I32Attr>:$space);
+  let assemblyFormat = [{
+    ` ` `<` `id` `=` $id `,` `mode` `=` $mode
+            (`,` `lbound` `=` $lbound^ `,` `ubound` `=` $ubound
+             `,` `space` `=` $space)?
+    `>` attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
 #endif // DXSA_OPS

mlir/lib/Dialect/DXSA/IR/DXSA.cpp

@@ -93,6 +93,15 @@ LogicalResult DclConstantBuffer::verify() {
   return success();
 }
 
+LogicalResult DclSampler::verify() {
+  auto lbound = getLbound();
+  auto ubound = getUbound();
+  if (lbound && ubound && *lbound > *ubound)
+    return emitOpError("expected lbound <= ubound, got lbound=")
+           << *lbound << ", ubound=" << *ubound;
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd attribute method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -704,6 +704,18 @@ class DXBuilder {
         optionalToAttr(space), accessPattern);
   }
 
+  Instruction buildDclSampler(uint32_t id, std::optional<uint32_t> lbound,
+                              std::optional<uint32_t> ubound,
+                              std::optional<uint32_t> space,
+                              dxsa::SamplerMode mode, Location loc) {
+    auto optionalToAttr = [&](std::optional<uint32_t> v) -> IntegerAttr {
+      return v ? builder.getI32IntegerAttr(*v) : IntegerAttr();
+    };
+    return dxsa::DclSampler::create(
+        builder, loc, id, mode, optionalToAttr(lbound), optionalToAttr(ubound),
+        optionalToAttr(space));
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -1386,6 +1398,35 @@ class Parser {
     }
   }
 
+  FailureOr<Instruction> parseDclSampler(uint32_t opcodeToken, Location loc) {
+    auto rawMode = DECODE_D3D10_SB_SAMPLER_MODE(opcodeToken);
+    auto mode = dxsa::symbolizeSamplerMode(rawMode);
+    if (!mode)
+      return emitError(loc, "unknown sampler mode: ") << rawMode;
+
+    auto operand = parseInlineOperand();
+    FAILURE_IF_FAILED(operand);
+    if (operand->getType() != dxsa::InlineOperandType::sampler)
+      return emitError(loc, "operand must be a sampler register, got ")
+             << dxsa::stringifyInlineOperandType(operand->getType());
+    auto indexArray = operand->getIndex();
+    auto indexDim = indexArray ? indexArray.size() : 0;
+    if (indexDim != 1 && indexDim != 3)
+      return emitError(loc, "operand must have a 1D or 3D index, got ")
+             << indexDim;
+    auto id = indexArray[0];
+    std::optional<uint32_t> lbound, ubound, space;
+    if (indexDim == 3) {
+      lbound = indexArray[1];
+      ubound = indexArray[2];
+      auto spaceToken = parseToken();
+      FAILURE_IF_FAILED(spaceToken);
+      space = *spaceToken;
+    }
+
+    return builder.buildDclSampler(id, lbound, ubound, space, *mode, loc);
+  }
+
   OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                           Instruction &out) {
     FailureOr<Instruction> result;
@@ -1465,6 +1506,9 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_CONSTANT_BUFFER:
       result = parseDclConstantBuffer(opcodeToken, loc);
       break;
+    case D3D10_SB_OPCODE_DCL_SAMPLER:
+      result = parseDclSampler(opcodeToken, loc);
+      break;
     default:
       return std::nullopt;
     }

mlir/test/Target/DXSA/dcl_sampler.mlir

@@ -0,0 +1,8 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_sampler.bin | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   dxsa.dcl_sampler <id = 0, mode = default>
+// CHECK-NEXT:   dxsa.dcl_sampler <id = 1, mode = comparison>
+// CHECK-NEXT:   dxsa.dcl_sampler <id = 2, mode = mono>
+// CHECK-NEXT:   dxsa.dcl_sampler <id = 0, mode = default, lbound = 0, ubound = 3, space = 1>
+// CHECK-NEXT: }

mlir/test/Target/DXSA/dcl_sampler_invalid.mlir

@@ -0,0 +1,4 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{expected lbound <= ubound, got lbound=5, ubound=3}}
+dxsa.dcl_sampler <id = 0, mode = default, lbound = 5, ubound = 3, space = 1>