[feat] Add bfloat16 value type support to the HKV for being enhanced by Ampere GPU BF16 training feature.

MoFHeka · rhdong · commit 42ad7d251e33 · 2024-02-05T11:30:06.000-08:00
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/hkv_hashtable_op_gpu.cu.cc b/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/hkv_hashtable_op_gpu.cu.cc
@@ -1048,6 +1048,7 @@ REGISTER_KERNEL(int64, int8);
 REGISTER_KERNEL(int64, int32);
 REGISTER_KERNEL(int64, int64);
 REGISTER_KERNEL(int64, Eigen::half);
+REGISTER_KERNEL(int64, Eigen::bfloat16);
 
 #undef REGISTER_KERNEL
 
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv_impl.cu.cc b/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv_impl.cu.cc
@@ -29,6 +29,7 @@ DEFINE_PURE_GPU_HASHTABLE(int64, int8);
 DEFINE_PURE_GPU_HASHTABLE(int64, int32);
 DEFINE_PURE_GPU_HASHTABLE(int64, int64);
 DEFINE_PURE_GPU_HASHTABLE(int64, Eigen::half);
+DEFINE_PURE_GPU_HASHTABLE(int64, Eigen::bfloat16);
 
 #undef DEFINE_PURE_GPU_HASHTABLE
 
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py
@@ -384,7 +384,7 @@ def test_variable(self):
       dim_list = [1, 2, 4, 8, 10, 16, 32, 64, 100, 200]
       kv_list = [[dtypes.int64, dtypes.float32], [dtypes.int64, dtypes.int32],
                  [dtypes.int64, dtypes.half], [dtypes.int64, dtypes.int8],
-                 [dtypes.int64, dtypes.int64]]
+                 [dtypes.int64, dtypes.int64], [dtypes.int64, dtypes.bfloat16]]
     else:
       dim_list = [1, 8, 16, 128]
       kv_list = [[dtypes.int32, dtypes.double], [dtypes.int32, dtypes.float32],
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py b/tensorflow_recommenders_addons/dynamic_embedding/python/ops/dynamic_embedding_variable.py
@@ -589,6 +589,7 @@ def _get_default_devices():
           [dtypes.int64, dtypes.int32],
           [dtypes.int64, dtypes.int64],
           [dtypes.int64, dtypes.half],
+          [dtypes.int64, dtypes.bfloat16],
       ]
     if is_macos() and is_arm64():
       if value_dtype == dtypes.half or value_dtype == dtypes.bfloat16:

Original file line number	Diff line number	Diff line change
`@@ -589,6 +589,7 @@ def _get_default_devices():`
`589`	`589`	`[dtypes.int64, dtypes.int32],`
`590`	`590`	`[dtypes.int64, dtypes.int64],`
`591`	`591`	`[dtypes.int64, dtypes.half],`
	`592`	`+ [dtypes.int64, dtypes.bfloat16],`
`592`	`593`	`]`
`593`	`594`	`if is_macos() and is_arm64():`
`594`	`595`	`if value_dtype == dtypes.half or value_dtype == dtypes.bfloat16:`