[mlir][dxsa] Add BinaryWriter to translate from MLIR to DXSA

BinaryWriter translates from an MLIR module in DXSA dialect into a
DXSA binary. It is a reverse of what BinaryParser does.

Current implementation only supports standard instructions, and needs
to be extended to support custom instructions.

Instruction table is moved into a separate file (InstrInfo.def), so it
can be shared between Parser/Writer, which build different data
structures from it. Parser goes from opcodes to mnemonics, and Writer
is reversed.

Tests are extended to run MLIR in roundtrip to verify both the Parser
and Writer. We also compare binary output with input to make sure that
we do not lose any data during translation.

Author: asavonic

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

@@ -128,7 +128,7 @@ def DXSA_IndexRelImm : DXSA_Op<"index.rel.imm"> {
     TODO
   }];
 
-  let arguments = (ins DXSA_OperandType:$operand, StrAttr:$op, I64Attr:$imm);
+  let arguments = (ins DXSA_OperandType:$operand, StrAttr:$op, I32Attr:$imm);
   let results = (outs DXSA_IndexType:$index);
   let assemblyFormat = "$operand attr-dict";
 }

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -17,10 +17,359 @@
 using namespace mlir;
 using namespace llvm;
 
+using OpcodeMap = llvm::DenseMap<StringRef, uint32_t>;
+
+static void initOpcodeMap(OpcodeMap &opcodes) {
+#define SET(OpCode, Name, NumOperands, PrecMask, OpClass)                      \
+  opcodes[Name] = OpCode;
+#include "InstrInfo.def"
+#undef SET
+}
+
+static FailureOr<uint32_t> getIndexRepresentation(Operation *op) {
+  if (auto imm = dyn_cast<dxsa::IndexImm>(op)) {
+    auto attr = dyn_cast<IntegerAttr>(imm.getImm());
+    if (!attr) {
+      emitError(op->getLoc(), "invalid immediate index");
+      return failure();
+    }
+
+    if (attr.getType().isInteger(32)) {
+      return D3D10_SB_OPERAND_INDEX_IMMEDIATE32;
+    }
+
+    if (attr.getType().isInteger(64)) {
+      return D3D10_SB_OPERAND_INDEX_IMMEDIATE64;
+    }
+
+    emitError(op->getLoc(), "invalid immediate index type");
+    return failure();
+  }
+
+  if (isa<dxsa::IndexRel>(op)) {
+    return D3D10_SB_OPERAND_INDEX_RELATIVE;
+  }
+
+  if (isa<dxsa::IndexRelImm>(op)) {
+    return D3D10_SB_OPERAND_INDEX_IMMEDIATE32_PLUS_RELATIVE;
+  }
+
+  emitError(op->getLoc(), "invalid index type");
+  return failure();
+}
+
+class Writer {
+public:
+  Writer(raw_ostream &output) : output(output, endianness::little) {
+    initOpcodeMap(opcodeMap);
+  }
+
+  LogicalResult emitModule(ModuleOp source) {
+    Region &region = source.getRegion();
+    if (!region.hasOneBlock()) {
+      emitError(region.getLoc(), "region should contain only one block");
+      return failure();
+    }
+
+    for (auto &op : region.front()) {
+      if (auto inst = dyn_cast<dxsa::Instruction>(op)) {
+        if (failed(emitInstruction(inst))) {
+          return failure();
+        }
+      }
+    }
+    return success();
+  }
+
+  // Emit an instruction and all its operands recursively.
+  // FIXME: add extended instructions
+  LogicalResult emitInstruction(dxsa::Instruction inst) {
+    // Buffer all tokens for an instruction, so we can fixup
+    // instruction length before emitting tokens to the output.
+    buffer.clear();
+
+    auto opcodeIt = opcodeMap.find(inst.getMnemonic());
+    if (opcodeIt == opcodeMap.end()) {
+      emitError(inst.getLoc(), "unknown mnemonic");
+      return failure();
+    }
+
+    // First token is an opcode and length. Length is unknown until we
+    // process all operands.
+    uint32_t opcode = opcodeIt->second;
+    uint32_t token = ENCODE_D3D10_SB_OPCODE_TYPE(opcode);
+    buffer.push_back(token);
+
+    for (Value value : inst.getOperands()) {
+      Operation *op = value.getDefiningOp();
+      if (!op) {
+        emitError(value.getLoc(), "undefined operand");
+        return failure();
+      }
+
+      if (auto operand = dyn_cast<dxsa::Operand>(*op)) {
+        if (failed(emitOperand(operand))) {
+          return failure();
+        }
+        continue;
+      }
+
+      if (auto operand = dyn_cast<dxsa::OperandImm>(*op)) {
+        if (failed(emitOperandImm(operand))) {
+          return failure();
+        }
+        continue;
+      }
+
+      emitError(op->getLoc(), "unexpected operand kind");
+      return failure();
+    }
+
+    // Fixup instruction length after all operands are accumulated in
+    // the buffer.
+    buffer[0] |= ENCODE_D3D10_SB_TOKENIZED_INSTRUCTION_LENGTH(buffer.size());
+    for (uint32_t token : buffer) {
+      output.write(token);
+    }
+
+    return success();
+  }
+
+  // Emit an operand and all its indices recursively.
+  LogicalResult emitOperand(dxsa::Operand op) {
+    uint32_t token = ENCODE_D3D10_SB_OPERAND_TYPE(op.getType());
+
+    // Encode swizzle, mask, or one component selection.
+    switch (op.getNumComponents()) {
+    case 0: {
+      token |=
+          ENCODE_D3D10_SB_OPERAND_NUM_COMPONENTS(D3D10_SB_OPERAND_0_COMPONENT);
+      break;
+    }
+    case 1: {
+      token |=
+          ENCODE_D3D10_SB_OPERAND_NUM_COMPONENTS(D3D10_SB_OPERAND_1_COMPONENT);
+      break;
+    }
+    case 4: {
+      token |=
+          ENCODE_D3D10_SB_OPERAND_NUM_COMPONENTS(D3D10_SB_OPERAND_4_COMPONENT);
+      if (auto mask = op.getMask()) {
+        token |= ENCODE_D3D10_SB_OPERAND_4_COMPONENT_SELECTION_MODE(
+            D3D10_SB_OPERAND_4_COMPONENT_MASK_MODE);
+        token |= ENCODE_D3D10_SB_OPERAND_4_COMPONENT_MASK(*mask);
+      } else if (auto swizzle = op.getSwizzle()) {
+        SmallVector<uint32_t, 4> values;
+        for (APInt v : *swizzle) {
+          values.push_back(v.getZExtValue());
+        }
+        if (values.size() != 4) {
+          emitError(op.getLoc(), "invalid number of swizzle values");
+          return failure();
+        }
+        token |= ENCODE_D3D10_SB_OPERAND_4_COMPONENT_SELECTION_MODE(
+            D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE_MODE);
+        token |= ENCODE_D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE(
+            values[0], values[1], values[2], values[3]);
+        break;
+      } else if (auto one = op.getOne()) {
+        token |= ENCODE_D3D10_SB_OPERAND_4_COMPONENT_SELECTION_MODE(
+            D3D10_SB_OPERAND_4_COMPONENT_SELECT_1_MODE);
+        token |= ENCODE_D3D10_SB_OPERAND_4_COMPONENT_SELECT_1(*one);
+        break;
+      }
+      break;
+    }
+    default: {
+      emitError(op.getLoc(), "invalid number of components");
+      return failure();
+    }
+    }
+
+    // Operand token encodes types and number of indices that follow
+    // it.
+    token |= ENCODE_D3D10_SB_OPERAND_INDEX_DIMENSION(op.getNumOperands());
+    uint32_t dim = 0;
+    for (Value value : op.getOperands()) {
+      Operation *index = value.getDefiningOp();
+      if (!index) {
+        emitError(value.getLoc(), "index must be defined");
+        return failure();
+      }
+
+      FailureOr<uint32_t> repr = getIndexRepresentation(index);
+      if (failed(repr)) {
+        return failure();
+      }
+      token |= ENCODE_D3D10_SB_OPERAND_INDEX_REPRESENTATION(dim, *repr);
+      dim += 1;
+    }
+
+    buffer.push_back(token);
+
+    // Indices follow the operand token.
+    for (Value value : op.getOperands()) {
+      Operation *index = value.getDefiningOp();
+      if (!index) {
+        emitError(value.getLoc(), "index must be defined");
+        return failure();
+      }
+
+      if (auto indexImm = dyn_cast<dxsa::IndexImm>(*index)) {
+        if (failed(emitIndexImm(indexImm))) {
+          return failure();
+        }
+        continue;
+      }
+
+      if (auto indexRel = dyn_cast<dxsa::IndexRel>(*index)) {
+        if (failed(emitIndexRel(indexRel))) {
+          return failure();
+        }
+        continue;
+      }
+
+      if (auto indexRelImm = dyn_cast<dxsa::IndexRelImm>(*index)) {
+        if (failed(emitIndexRelImm(indexRelImm))) {
+          return failure();
+        }
+        continue;
+      }
+
+      emitError(value.getLoc(), "invalid index type");
+      return failure();
+    }
+
+    return success();
+  }
+
+  // Emit an immediate operand. Unlike register operands, immediate
+  // operands do not have indices. They are encoded as an operand
+  // followed by N immediate values for each component.
+  LogicalResult emitOperandImm(dxsa::OperandImm op) {
+    auto attr = dyn_cast<DenseIntElementsAttr>(op.getImm());
+    if (!attr) {
+      emitError(op.getLoc(), "invalid immediate operand");
+    }
+
+    uint32_t token = 0;
+
+    Type elementType = attr.getType().getElementType();
+    if (elementType.isInteger(32)) {
+      token |= ENCODE_D3D10_SB_OPERAND_TYPE(D3D10_SB_OPERAND_TYPE_IMMEDIATE32);
+    } else if (elementType.isInteger(64)) {
+      token |= ENCODE_D3D10_SB_OPERAND_TYPE(D3D10_SB_OPERAND_TYPE_IMMEDIATE64);
+    } else {
+      emitError(op.getLoc(), "invalid immediate operand type");
+      return failure();
+    }
+
+    // Split immediates into tokens. 32 bit immediate values are
+    // encoded as is, and 64 bit immediates are split into high and
+    // low 32 bit parts.
+    SmallVector<uint32_t, 4> values;
+    for (APInt v : attr) {
+      uint64_t bits = v.getZExtValue();
+      if (v.getBitWidth() == 64) {
+        values.push_back(bits >> 32);
+      }
+      values.push_back(bits);
+    }
+
+    if (values.size() == 1) {
+      token |=
+          ENCODE_D3D10_SB_OPERAND_NUM_COMPONENTS(D3D10_SB_OPERAND_1_COMPONENT);
+    } else if (values.size() == 4) {
+      token |=
+          ENCODE_D3D10_SB_OPERAND_NUM_COMPONENTS(D3D10_SB_OPERAND_4_COMPONENT);
+    } else {
+      emitError(op.getLoc(),
+                "immediate operand should be either 1- or 4- component");
+      return failure();
+    }
+
+    buffer.push_back(token);
+    for (uint32_t v : values) {
+      buffer.push_back(v);
+    }
+
+    return success();
+  }
+
+  // Emit an immediate index. Its type is encoded into the operand, so
+  // here we only emit the value as tokens.
+  LogicalResult emitIndexImm(dxsa::IndexImm op) {
+    auto attr = dyn_cast<IntegerAttr>(op.getImm());
+    if (!attr) {
+      emitError(op.getLoc(), "invalid immediate index");
+      return failure();
+    }
+
+    uint64_t value = attr.getInt();
+    if (attr.getType().isInteger(32)) {
+      buffer.push_back(value);
+      return success();
+    }
+
+    if (attr.getType().isInteger(64)) {
+      buffer.push_back(value >> 32);
+      buffer.push_back(value);
+      return success();
+    }
+
+    emitError(op.getLoc(), "invalid type of an immediate index");
+    return failure();
+  }
+
+  // Emit an operand used as an index.
+  LogicalResult emitIndexRel(dxsa::IndexRel index) {
+    Operation *def = index.getOperand().getDefiningOp();
+    if (!def) {
+      emitError(index.getLoc(), "index must be defined");
+      return failure();
+    }
+
+    auto operand = dyn_cast<dxsa::Operand>(*def);
+    if (!operand) {
+      emitError(def->getLoc(), "invalid index relative operand");
+      return failure();
+    }
+
+    // Recursively emit an operand, which may also have other indices.
+    return emitOperand(operand);
+  }
+
+  // Emit an index as an operand + a 32 bit immediate offset.
+  LogicalResult emitIndexRelImm(dxsa::IndexRelImm index) {
+    Operation *def = index.getOperand().getDefiningOp();
+    if (!def) {
+      emitError(index.getLoc(), "index must be defined");
+      return failure();
+    }
+
+    auto operand = dyn_cast<dxsa::Operand>(*def);
+    if (!operand) {
+      emitError(def->getLoc(), "invalid index relative operand");
+      return failure();
+    }
+
+    if (failed(emitOperand(operand))) {
+      return failure();
+    }
+
+    buffer.push_back(index.getImm());
+    return success();
+  }
+
+private:
+  std::vector<uint32_t> buffer;
+  support::endian::Writer output;
+  OpcodeMap opcodeMap;
+};
+
 namespace mlir::dxsa {
 LogicalResult exportModuleToDxsaBinary(ModuleOp source, raw_ostream &output) {
-  Region &region = source.getRegion();
-  assert(region.hasOneBlock() && "invalid module");
-  return failure();
+  Writer writer(output);
+  return writer.emitModule(source);
 }
 } // namespace mlir::dxsa