Merge pull request #17307 from vespa-engine/lesters/new-tensor-cell-types-java

Add bfloat16 and int8 tensor cell types in Java

2 files changed, 42 insertions, 11 deletions

diff --git a/eval/src/apps/make_tensor_binary_format_test_spec/make_tensor_binary_format_test_spec.cpp b/eval/src/apps/make_tensor_binary_format_test_spec/make_tensor_binary_format_test_spec.cpp
index 6e882fc3d9d..974f95a2add 100644
--- a/eval/src/apps/make_tensor_binary_format_test_spec/make_tensor_binary_format_test_spec.cpp
+++ b/eval/src/apps/make_tensor_binary_format_test_spec/make_tensor_binary_format_test_spec.cpp
@@ -3,6 +3,8 @@
 #include <vespa/vespalib/data/slime/slime.h>
 #include <vespa/vespalib/objects/nbostream.h>
 #include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/vespalib/util/bfloat16.h>
+#include <vespa/eval/eval/int8float.h>
 #include <vespa/eval/eval/tensor_spec.h>
 #include <vespa/eval/eval/value_type.h>
 #include <vespa/eval/eval/test/test_io.h>
@@ -20,14 +22,20 @@ using Dict = std::vector<vespalib::string>;
 template <typename T> std::vector<bool> with_cell_type_opts();
 template <> std::vector<bool> with_cell_type_opts<double>() { return {false, true}; }
 template <> std::vector<bool> with_cell_type_opts<float>() { return {true}; }
+template <> std::vector<bool> with_cell_type_opts<BFloat16>() { return {true}; }
+template <> std::vector<bool> with_cell_type_opts<Int8Float>() { return {true}; }
 
 template <typename T> uint8_t cell_type_id();
 template <> uint8_t cell_type_id<double>() { return 0; }
 template <> uint8_t cell_type_id<float>() { return 1; }
+template <> uint8_t cell_type_id<BFloat16>() { return 2; }
+template <> uint8_t cell_type_id<Int8Float>() { return 3; }
 
 template <typename T> const char *cell_type_str();
 template <> const char *cell_type_str<double>() { return ""; }
 template <> const char *cell_type_str<float>() { return "<float>"; }
+template <> const char *cell_type_str<BFloat16>() { return "<bfloat16>"; }
+template <> const char *cell_type_str<Int8Float>() { return "<int8>"; }
 
 template <typename T> nbostream make_sparse(bool with_cell_type) {
     nbostream data;
@@ -62,7 +70,8 @@ template <typename T> nbostream make_mixed(bool with_cell_type) {
     return data;
 }
 
-void set_tensor(Cursor &test, const TensorSpec &spec) {
+void set_tensor(Cursor &test, const TensorSpec &spec_in) {
+    auto spec = spec_in.normalize();
     const Inspector &old_tensor = test["tensor"];
     if (old_tensor.valid()) {
         TensorSpec old_spec = TensorSpec::from_slime(old_tensor);
@@ -183,8 +192,8 @@ void make_vector_test(Cursor &test, size_t x_size) {
         for (size_t x = 0; x < x_size; ++x) {
             double value = val(x);
             spec.add({{"x", x}}, value);
-            dense << static_cast<T>(value);
-            mixed << static_cast<T>(value);
+            dense << T(value);
+            mixed << T(value);
         }
         set_tensor(test, spec);
         add_binary(test, {dense, mixed});
@@ -212,8 +221,8 @@ void make_matrix_test(Cursor &test, size_t x_size, size_t y_size) {
             for (size_t y = 0; y < y_size; ++y) {
                 double value = mix({val(x), val(y)});
                 spec.add({{"x", x}, {"y", y}}, value);
-                dense << static_cast<T>(value);
-                mixed << static_cast<T>(value);
+                dense << T(value);
+                mixed << T(value);
             }
         }
         set_tensor(test, spec);
@@ -245,8 +254,8 @@ void make_map_test(Cursor &test, const Dict &x_dict_in) {
                 spec.add({{"x", x}}, value);
                 sparse.writeSmallString(x);
                 mixed.writeSmallString(x);
-                sparse << static_cast<T>(value);
-                mixed << static_cast<T>(value);
+                sparse << T(value);
+                mixed << T(value);
             }
             set_tensor(test, spec);
             add_binary(test, {sparse, mixed});
@@ -285,8 +294,8 @@ void make_mesh_test(Cursor &test, const Dict &x_dict_in, const vespalib::string
                 sparse.writeSmallString(y);
                 mixed.writeSmallString(x);
                 mixed.writeSmallString(y);
-                sparse << static_cast<T>(value);
-                mixed << static_cast<T>(value);
+                sparse << T(value);
+                mixed << T(value);
             }
             set_tensor(test, spec);
             add_binary(test, {sparse, mixed});
@@ -326,7 +335,7 @@ void make_vector_map_test(Cursor &test,
                 for (size_t idx = 0; idx < indexed_size; ++idx) {
                     double value = mix({val(label), val(idx)});
                     spec.add({{mapped_name, label}, {indexed_name, idx}}, value);
-                    mixed << static_cast<T>(value);
+                    mixed << T(value);
                 }
             }
             set_tensor(test, spec);
@@ -360,6 +369,8 @@ void make_tests(test::TestWriter &writer) {
     make_number_test(writer.create(), 42.0);
     make_typed_tests<double>(writer);
     make_typed_tests<float>(writer);
+    make_typed_tests<BFloat16>(writer);
+    make_typed_tests<Int8Float>(writer);
 }
 
 int main(int, char **) {
diff --git a/eval/src/apps/make_tensor_binary_format_test_spec/test_spec.json b/eval/src/apps/make_tensor_binary_format_test_spec/test_spec.json
index f6b535e071a..b7710eadf5d 100644
--- a/eval/src/apps/make_tensor_binary_format_test_spec/test_spec.json
+++ b/eval/src/apps/make_tensor_binary_format_test_spec/test_spec.json
@@ -20,4 +20,24 @@
 {"tensor":{"type":"tensor<float>(x[10],y{})","cells":[]},"binary":["0x07010101790101780A00"]}
 {"tensor":{"type":"tensor<float>(x{},y[3])","cells":[{"address":{"x":"a","y":0},"value":11},{"address":{"x":"a","y":1},"value":12},{"address":{"x":"a","y":2},"value":13},{"address":{"x":"b","y":0},"value":21},{"address":{"x":"b","y":1},"value":22},{"address":{"x":"b","y":2},"value":23}]},"binary":["0x070101017801017903020161413000004140000041500000016241A8000041B0000041B80000","0x07010101780101790302016241A8000041B0000041B800000161413000004140000041500000"]}
 {"tensor":{"type":"tensor<float>(x[3],y{})","cells":[{"address":{"x":0,"y":"a"},"value":11},{"address":{"x":0,"y":"b"},"value":21},{"address":{"x":1,"y":"a"},"value":12},{"address":{"x":1,"y":"b"},"value":22},{"address":{"x":2,"y":"a"},"value":13},{"address":{"x":2,"y":"b"},"value":23}]},"binary":["0x070101017901017803020161413000004140000041500000016241A8000041B0000041B80000","0x07010101790101780302016241A8000041B0000041B800000161413000004140000041500000"]}
-{"num_tests":22}
+{"tensor":{"type":"tensor<bfloat16>(x[3])","cells":[{"address":{"x":0},"value":1},{"address":{"x":1},"value":2},{"address":{"x":2},"value":3}]},"binary":["0x0602010178033F8040004040","0x070200010178033F8040004040"]}
+{"tensor":{"type":"tensor<bfloat16>(x[2],y[3])","cells":[{"address":{"x":0,"y":0},"value":11},{"address":{"x":0,"y":1},"value":12},{"address":{"x":0,"y":2},"value":13},{"address":{"x":1,"y":0},"value":21},{"address":{"x":1,"y":1},"value":22},{"address":{"x":1,"y":2},"value":23}]},"binary":["0x06020201780201790341304140415041A841B041B8","0x0702000201780201790341304140415041A841B041B8"]}
+{"tensor":{"type":"tensor<bfloat16>(x{})","cells":[]},"binary":["0x050201017800","0x07020101780000"]}
+{"tensor":{"type":"tensor<bfloat16>(x{})","cells":[{"address":{"x":"a"},"value":1},{"address":{"x":"b"},"value":2},{"address":{"x":"c"},"value":3}]},"binary":["0x05020101780301613F800162400001634040","0x0702010178000301613F800162400001634040","0x05020101780301613F800163404001624000","0x0702010178000301613F800163404001624000","0x0502010178030162400001613F8001634040","0x070201017800030162400001613F8001634040","0x050201017803016240000163404001613F80","0x07020101780003016240000163404001613F80","0x0502010178030163404001613F8001624000","0x070201017800030163404001613F8001624000","0x050201017803016340400162400001613F80","0x07020101780003016340400162400001613F80"]}
+{"tensor":{"type":"tensor<bfloat16>(x{},y{})","cells":[]},"binary":["0x0502020178017900","0x070202017801790000"]}
+{"tensor":{"type":"tensor<bfloat16>(x{},y{})","cells":[{"address":{"x":"bar","y":"a"},"value":21},{"address":{"x":"foo","y":"a"},"value":11}]},"binary":["0x050202017801790203666F6F0161413003626172016141A8","0x07020201780179000203666F6F0161413003626172016141A8","0x050202017801790203626172016141A803666F6F01614130","0x07020201780179000203626172016141A803666F6F01614130"]}
+{"tensor":{"type":"tensor<bfloat16>(x{},y[10])","cells":[]},"binary":["0x07020101780101790A00"]}
+{"tensor":{"type":"tensor<bfloat16>(x[10],y{})","cells":[]},"binary":["0x07020101790101780A00"]}
+{"tensor":{"type":"tensor<bfloat16>(x{},y[3])","cells":[{"address":{"x":"a","y":0},"value":11},{"address":{"x":"a","y":1},"value":12},{"address":{"x":"a","y":2},"value":13},{"address":{"x":"b","y":0},"value":21},{"address":{"x":"b","y":1},"value":22},{"address":{"x":"b","y":2},"value":23}]},"binary":["0x070201017801017903020161413041404150016241A841B041B8","0x07020101780101790302016241A841B041B80161413041404150"]}
+{"tensor":{"type":"tensor<bfloat16>(x[3],y{})","cells":[{"address":{"x":0,"y":"a"},"value":11},{"address":{"x":0,"y":"b"},"value":21},{"address":{"x":1,"y":"a"},"value":12},{"address":{"x":1,"y":"b"},"value":22},{"address":{"x":2,"y":"a"},"value":13},{"address":{"x":2,"y":"b"},"value":23}]},"binary":["0x070201017901017803020161413041404150016241A841B041B8","0x07020101790101780302016241A841B041B80161413041404150"]}
+{"tensor":{"type":"tensor<int8>(x[3])","cells":[{"address":{"x":0},"value":1},{"address":{"x":1},"value":2},{"address":{"x":2},"value":3}]},"binary":["0x060301017803010203","0x07030001017803010203"]}
+{"tensor":{"type":"tensor<int8>(x[2],y[3])","cells":[{"address":{"x":0,"y":0},"value":11},{"address":{"x":0,"y":1},"value":12},{"address":{"x":0,"y":2},"value":13},{"address":{"x":1,"y":0},"value":21},{"address":{"x":1,"y":1},"value":22},{"address":{"x":1,"y":2},"value":23}]},"binary":["0x0603020178020179030B0C0D151617","0x070300020178020179030B0C0D151617"]}
+{"tensor":{"type":"tensor<int8>(x{})","cells":[]},"binary":["0x050301017800","0x07030101780000"]}
+{"tensor":{"type":"tensor<int8>(x{})","cells":[{"address":{"x":"a"},"value":1},{"address":{"x":"b"},"value":2},{"address":{"x":"c"},"value":3}]},"binary":["0x050301017803016101016202016303","0x07030101780003016101016202016303","0x050301017803016101016303016202","0x07030101780003016101016303016202","0x050301017803016202016101016303","0x07030101780003016202016101016303","0x050301017803016202016303016101","0x07030101780003016202016303016101","0x050301017803016303016101016202","0x07030101780003016303016101016202","0x050301017803016303016202016101","0x07030101780003016303016202016101"]}
+{"tensor":{"type":"tensor<int8>(x{},y{})","cells":[]},"binary":["0x0503020178017900","0x070302017801790000"]}
+{"tensor":{"type":"tensor<int8>(x{},y{})","cells":[{"address":{"x":"bar","y":"a"},"value":21},{"address":{"x":"foo","y":"a"},"value":11}]},"binary":["0x050302017801790203666F6F01610B03626172016115","0x07030201780179000203666F6F01610B03626172016115","0x05030201780179020362617201611503666F6F01610B","0x0703020178017900020362617201611503666F6F01610B"]}
+{"tensor":{"type":"tensor<int8>(x{},y[10])","cells":[]},"binary":["0x07030101780101790A00"]}
+{"tensor":{"type":"tensor<int8>(x[10],y{})","cells":[]},"binary":["0x07030101790101780A00"]}
+{"tensor":{"type":"tensor<int8>(x{},y[3])","cells":[{"address":{"x":"a","y":0},"value":11},{"address":{"x":"a","y":1},"value":12},{"address":{"x":"a","y":2},"value":13},{"address":{"x":"b","y":0},"value":21},{"address":{"x":"b","y":1},"value":22},{"address":{"x":"b","y":2},"value":23}]},"binary":["0x0703010178010179030201610B0C0D0162151617","0x07030101780101790302016215161701610B0C0D"]}
+{"tensor":{"type":"tensor<int8>(x[3],y{})","cells":[{"address":{"x":0,"y":"a"},"value":11},{"address":{"x":0,"y":"b"},"value":21},{"address":{"x":1,"y":"a"},"value":12},{"address":{"x":1,"y":"b"},"value":22},{"address":{"x":2,"y":"a"},"value":13},{"address":{"x":2,"y":"b"},"value":23}]},"binary":["0x0703010179010178030201610B0C0D0162151617","0x07030101790101780302016215161701610B0C0D"]}
+{"num_tests":42}