Get started on bfloat16 support

Author: Bisaloo

R/codecs_array_bytes.R

@@ -5,6 +5,16 @@ codec_bytes_decode <- function(
   datatype,
   endian
 ) {
+  if (!is.list(datatype) && datatype$base_type == "bfloat") {
+    return(.Call(
+      "type_convert_bfloat",
+      input,
+      datatype$nbytes,
+      chunk_dim,
+      endian,
+      PACKAGE = "Rarr"
+    ))
+  }
   convert_bytes_to_array(
     input,
     datatype$base_type,

src/type_conversion.c

@@ -0,0 +1,45 @@
+#include "type_conversion.h"
+
+SEXP type_convert_bfloat(SEXP input, SEXP _n_bytes, SEXP dims, SEXP endian) {
+
+  const int n_bytes = INTEGER(_n_bytes)[0];
+  const R_xlen_t length = xlength(input);
+  const bool big_endian = strcmp(CHAR(STRING_ELT(endian, 0)), "big") == 0 ? 1 : 0;
+  const uint8_t* raw_buffer = (const uint8_t*) RAW(input);
+
+  SEXP data;
+  const R_xlen_t data_length = length / n_bytes;
+  R_xlen_t i;
+
+  // space for the converted output
+  data = PROTECT(allocVector(REALSXP, data_length));
+  double *p_data = REAL(data);
+
+  if(n_bytes != 2) {
+    Rf_error("Only 2 byte bfloat16 is supported");
+  }
+
+  for (i = 0; i < data_length; i++) {
+    uint8_t b0 = raw_buffer[i * 2];
+    uint8_t b1 = raw_buffer[i * 2 + 1];
+    uint32_t f32_bits;
+    if (big_endian) {
+      f32_bits = (uint32_t)b0 << 24 | (uint32_t)b1 << 16;
+    } else {
+      f32_bits = (uint32_t)b0 | (uint32_t)b1 << 8;
+      f32_bits <<= 16;
+    }
+    float f32_val;
+    memcpy(&f32_val, &f32_bits, sizeof(float));
+    // Magic scaling factor to round when converting bfloat16 to float32, 
+    // from https://stackoverflow.com/a/55290557
+    p_data[i] = (double)f32_val * 1.001957f;
+  }
+
+  if (!isNull(dims) && xlength(dims) > 0) {
+    Rf_dimgets(data, dims);
+  }
+  
+  UNPROTECT(1);
+  return(data);
+}

src/type_conversion.h

@@ -0,0 +1,3 @@
+#include "Rarr.h"
+
+SEXP type_convert_bfloat(SEXP input, SEXP _n_bytes, SEXP dims, SEXP endian);