Add initial support for SPV_EXT_float8 (KhronosGroup#6170)

* Add initial support for SPV_EXT_float8

This also overhauls the number parsing code. It seems further work
will be needed to correctly handle BFloat16. This change intends to
leave the code no worse than it was. TODOs were added where relevant.

Change-Id: I0e146a949ae68d0488dabbb91191f2801ba4b49c
Signed-off-by: Guillaume Trebuchet <[email protected]>
Signed-off-by: Kevin Petit <[email protected]>

* attempt to fix MSVC warning

Change-Id: I630635c54393d7ffc82001f2eec0d2538300521e

* another attempt

Change-Id: I237f6c8ca054347a8e5da62f7649f6fb1d76cfa1

* address review comments

Change-Id: Ic0be831943f040623860ccc4dc2321c3616f10d1

* use more descriptive FP encoding enum names for FP8

Change-Id: I8ba05748480a030e0246fba14996fb123b1f146b

* add a few more validator tests

Change-Id: Idfa90b67d69cabcd092c5cc1dbe06f553946c73a

* format fixes

Change-Id: I1e90212170eaa547019e6bd3672b64ccc2ab3871

---------

Signed-off-by: Guillaume Trebuchet <[email protected]>
Signed-off-by: Kevin Petit <[email protected]>

Author: kpet

DEPS

@@ -14,7 +14,7 @@ vars = {
 
   're2_revision': 'c84a140c93352cdabbfb547c531be34515b12228',
 
-  'spirv_headers_revision': '7168a5ad041f6b6b9170f027c7417f98a2056ff0',
+  'spirv_headers_revision': 'fd96661925488574fe247a779babe5d380b63635',
 }
 
 deps = {

include/spirv-tools/libspirv.h

@@ -390,6 +390,18 @@ typedef enum spv_number_kind_t {
   SPV_NUMBER_FLOATING,
 } spv_number_kind_t;
 
+// Represent the encoding of floating point values
+typedef enum spv_fp_encoding_t {
+  SPV_FP_ENCODING_UNKNOWN =
+      0,  // The encoding is not specified. Has to be deduced from bitwidth
+  SPV_FP_ENCODING_IEEE754_BINARY16,  // half float
+  SPV_FP_ENCODING_IEEE754_BINARY32,  // single float
+  SPV_FP_ENCODING_IEEE754_BINARY64,  // double float
+  SPV_FP_ENCODING_BFLOAT16,
+  SPV_FP_ENCODING_FLOAT8_E4M3,
+  SPV_FP_ENCODING_FLOAT8_E5M2,
+} spv_fp_encoding_t;
+
 typedef enum spv_text_to_binary_options_t {
   SPV_TEXT_TO_BINARY_OPTION_NONE = SPV_BIT(0),
   // Numeric IDs in the binary will have the same values as in the source.
@@ -445,6 +457,8 @@ typedef struct spv_parsed_operand_t {
   spv_number_kind_t number_kind;
   // The number of bits for a literal number type.
   uint32_t number_bit_width;
+  // The encoding used for floating point values
+  spv_fp_encoding_t fp_encoding;
 } spv_parsed_operand_t;
 
 // An instruction parsed from a binary SPIR-V module.

source/binary.cpp

@@ -190,6 +190,7 @@ class Parser {
   struct NumberType {
     spv_number_kind_t type;
     uint32_t bit_width;
+    spv_fp_encoding_t encoding;
   };
 
   // The state used to parse a single SPIR-V binary module.
@@ -385,8 +386,6 @@ spv_result_t Parser::parseInstruction() {
   assert(_.requires_endian_conversion ||
          (_.endian_converted_words.size() == 1));
 
-  recordNumberType(inst_offset, &inst);
-
   if (_.requires_endian_conversion) {
     // We must wait until here to set this pointer, because the vector might
     // have been be resized while we accumulated its elements.
@@ -398,6 +397,8 @@ spv_result_t Parser::parseInstruction() {
   }
   inst.num_words = inst_word_count;
 
+  recordNumberType(inst_offset, &inst);
+
   // We must wait until here to set this pointer, because the vector might
   // have been be resized while we accumulated its elements.
   inst.operands = _.operands.data();
@@ -833,6 +834,7 @@ spv_result_t Parser::setNumericTypeInfoForType(
 
   parsed_operand->number_kind = info.type;
   parsed_operand->number_bit_width = info.bit_width;
+  parsed_operand->fp_encoding = info.encoding;
   // Round up the word count.
   parsed_operand->num_words = static_cast<uint16_t>((info.bit_width + 31) / 32);
   return SPV_SUCCESS;
@@ -850,6 +852,17 @@ void Parser::recordNumberType(size_t inst_offset,
     } else if (spv::Op::OpTypeFloat == opcode) {
       info.type = SPV_NUMBER_FLOATING;
       info.bit_width = peekAt(inst_offset + 2);
+      if (inst->num_words >= 4) {
+        const spvtools::OperandDesc* desc;
+        spv_result_t status = spvtools::LookupOperand(
+            SPV_OPERAND_TYPE_FPENCODING, peekAt(inst_offset + 3), &desc);
+        if (status == SPV_SUCCESS) {
+          info.encoding = spvFPEncodingFromOperandFPEncoding(
+              static_cast<spv::FPEncoding>(desc->value));
+        } else {
+          info.encoding = SPV_FP_ENCODING_UNKNOWN;
+        }
+      }
     }
     // The *result* Id of a type generating instruction is the type Id.
     _.type_id_to_number_type_info[inst->result_id] = info;

source/name_mapper.cpp

@@ -217,6 +217,14 @@ spv_result_t FriendlyNameMapper::ParseInstruction(
           SaveName(result_id, "bfloat16");
           break;
         }
+        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float8E4M3EXT) {
+          SaveName(result_id, "fp8e4m3");
+          break;
+        }
+        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float8E5M2EXT) {
+          SaveName(result_id, "fp8e5m2");
+          break;
+        }
       }
       switch (bit_width) {
         case 16:

source/operand.cpp

@@ -619,3 +619,17 @@ std::function<bool(unsigned)> spvDbgInfoExtOperandCanBeForwardDeclaredFunction(
   }
   return out;
 }
+
+spv_fp_encoding_t spvFPEncodingFromOperandFPEncoding(spv::FPEncoding encoding) {
+  switch (encoding) {
+    case spv::FPEncoding::BFloat16KHR:
+      return SPV_FP_ENCODING_BFLOAT16;
+    case spv::FPEncoding::Float8E4M3EXT:
+      return SPV_FP_ENCODING_FLOAT8_E4M3;
+    case spv::FPEncoding::Float8E5M2EXT:
+      return SPV_FP_ENCODING_FLOAT8_E5M2;
+    case spv::FPEncoding::Max:
+      break;
+  }
+  return SPV_FP_ENCODING_UNKNOWN;
+}

source/operand.h

@@ -122,4 +122,7 @@ std::function<bool(unsigned)> spvOperandCanBeForwardDeclaredFunction(
 std::function<bool(unsigned)> spvDbgInfoExtOperandCanBeForwardDeclaredFunction(
     spv::Op opcode, spv_ext_inst_type_t ext_type, uint32_t key);
 
+// Converts an spv::FPEncoding to spv_fp_encoding_t
+spv_fp_encoding_t spvFPEncodingFromOperandFPEncoding(spv::FPEncoding encoding);
+
 #endif  // SOURCE_OPERAND_H_

source/opt/types.cpp

@@ -320,6 +320,14 @@ std::string Float::str() const {
       assert(width_ == 16);
       oss << "bfloat16";
       break;
+    case spv::FPEncoding::Float8E4M3EXT:
+      assert(width_ == 8);
+      oss << "fp8e4m3";
+      break;
+    case spv::FPEncoding::Float8E5M2EXT:
+      assert(width_ == 8);
+      oss << "fp8e5m2";
+      break;
     default:
       oss << "float" << width_;
       break;

source/parsed_operand.cpp

@@ -43,12 +43,35 @@ void EmitNumericLiteral(std::ostream* out, const spv_parsed_instruction_t& inst,
         *out << word;
         break;
       case SPV_NUMBER_FLOATING:
-        if (operand.number_bit_width == 16) {
-          *out << spvtools::utils::FloatProxy<spvtools::utils::Float16>(
-              uint16_t(word & 0xFFFF));
-        } else {
-          // Assume 32-bit floats.
-          *out << spvtools::utils::FloatProxy<float>(word);
+        switch (operand.fp_encoding) {
+          case SPV_FP_ENCODING_IEEE754_BINARY16:
+            *out << spvtools::utils::FloatProxy<spvtools::utils::Float16>(
+                uint16_t(word & 0xFFFF));
+            break;
+          case SPV_FP_ENCODING_IEEE754_BINARY32:
+            *out << spvtools::utils::FloatProxy<float>(word);
+            break;
+          case SPV_FP_ENCODING_FLOAT8_E4M3:
+            *out << spvtools::utils::FloatProxy<spvtools::utils::Float8_E4M3>(
+                uint8_t(word & 0xFF));
+            break;
+          case SPV_FP_ENCODING_FLOAT8_E5M2:
+            *out << spvtools::utils::FloatProxy<spvtools::utils::Float8_E5M2>(
+                uint8_t(word & 0xFF));
+            break;
+          // TODO Bfloat16
+          case SPV_FP_ENCODING_UNKNOWN:
+            switch (operand.number_bit_width) {
+              case 16:
+                *out << spvtools::utils::FloatProxy<spvtools::utils::Float16>(
+                    uint16_t(word & 0xFFFF));
+                break;
+              case 32:
+                *out << spvtools::utils::FloatProxy<float>(word);
+                break;
+            }
+          default:
+            break;
         }
         break;
       default:

source/text_handler.cpp

@@ -254,13 +254,13 @@ spv_result_t AssemblyContext::binaryEncodeNumericLiteral(
              << "Unexpected numeric literal type";
     case IdTypeClass::kScalarIntegerType:
       if (type.isSigned) {
-        number_type = {type.bitwidth, SPV_NUMBER_SIGNED_INT};
+        number_type = {type.bitwidth, SPV_NUMBER_SIGNED_INT, type.encoding};
       } else {
-        number_type = {type.bitwidth, SPV_NUMBER_UNSIGNED_INT};
+        number_type = {type.bitwidth, SPV_NUMBER_UNSIGNED_INT, type.encoding};
       }
       break;
     case IdTypeClass::kScalarFloatType:
-      number_type = {type.bitwidth, SPV_NUMBER_FLOATING};
+      number_type = {type.bitwidth, SPV_NUMBER_FLOATING, type.encoding};
       break;
     case IdTypeClass::kBottom:
       // kBottom means the type is unknown and we need to infer the type before
@@ -270,11 +270,11 @@ spv_result_t AssemblyContext::binaryEncodeNumericLiteral(
       // signed integer, otherwise an unsigned integer.
       uint32_t bitwidth = static_cast<uint32_t>(assumedBitWidth(type));
       if (strchr(val, '.')) {
-        number_type = {bitwidth, SPV_NUMBER_FLOATING};
+        number_type = {bitwidth, SPV_NUMBER_FLOATING, type.encoding};
       } else if (type.isSigned || val[0] == '-') {
-        number_type = {bitwidth, SPV_NUMBER_SIGNED_INT};
+        number_type = {bitwidth, SPV_NUMBER_SIGNED_INT, type.encoding};
       } else {
-        number_type = {bitwidth, SPV_NUMBER_UNSIGNED_INT};
+        number_type = {bitwidth, SPV_NUMBER_UNSIGNED_INT, type.encoding};
       }
       break;
   }
@@ -330,14 +330,27 @@ spv_result_t AssemblyContext::recordTypeDefinition(
     if (pInst->words.size() != 4)
       return diagnostic() << "Invalid OpTypeInt instruction";
     types_[value] = {pInst->words[2], pInst->words[3] != 0,
-                     IdTypeClass::kScalarIntegerType};
+                     IdTypeClass::kScalarIntegerType, SPV_FP_ENCODING_UNKNOWN};
   } else if (pInst->opcode == spv::Op::OpTypeFloat) {
     if ((pInst->words.size() != 3) && (pInst->words.size() != 4))
       return diagnostic() << "Invalid OpTypeFloat instruction";
-    // TODO(kpet) Do we need to record the FP Encoding here?
-    types_[value] = {pInst->words[2], false, IdTypeClass::kScalarFloatType};
+    spv_fp_encoding_t enc = SPV_FP_ENCODING_UNKNOWN;
+    if (pInst->words.size() >= 4) {
+      const spvtools::OperandDesc* desc;
+      spv_result_t status = spvtools::LookupOperand(SPV_OPERAND_TYPE_FPENCODING,
+                                                    pInst->words[3], &desc);
+      if (status == SPV_SUCCESS) {
+        enc = spvFPEncodingFromOperandFPEncoding(
+            static_cast<spv::FPEncoding>(desc->value));
+      } else {
+        return diagnostic() << "Invalid OpTypeFloat encoding";
+      }
+    }
+    types_[value] = {pInst->words[2], false, IdTypeClass::kScalarFloatType,
+                     enc};
   } else {
-    types_[value] = {0, false, IdTypeClass::kOtherType};
+    types_[value] = {0, false, IdTypeClass::kOtherType,
+                     SPV_FP_ENCODING_UNKNOWN};
   }
   return SPV_SUCCESS;
 }

source/text_handler.h

@@ -47,6 +47,7 @@ struct IdType {
   uint32_t bitwidth;  // Safe to assume that we will not have > 2^32 bits.
   bool isSigned;      // This is only significant if type_class is integral.
   IdTypeClass type_class;
+  spv_fp_encoding_t encoding;
 };
 
 // Default equality operator for IdType. Tests if all members are the same.