Merge pull request #7 from gomlx/isfinite

Added IsInfinite and PopulationCount operations.
gomlx · Oct 3, 2024 · c818cc3 · c818cc3
2 parents c0f3179 + c74d37b
commit c818cc3
Show file tree

Hide file tree

Showing 13 changed files with 1,987 additions and 1,716 deletions.
diff --git a/c/gomlx/xlabuilder/gen_op_types.h b/c/gomlx/xlabuilder/gen_op_types.h
@@ -90,6 +90,8 @@ enum OpType {
   DynamicSliceOp,
   DynamicUpdateSliceOp,
   ErfOp,
+  IsFiniteOp,
+  PopulationCountOp,
 };
 
 #ifdef __cplusplus

diff --git a/c/gomlx/xlabuilder/xlabuilder.cpp b/c/gomlx/xlabuilder/xlabuilder.cpp
@@ -532,6 +532,12 @@ XlaStatus *XlaBuilderAddOp(XlaBuilder *builder, SerializedOp *serialized_op) {
   case ErfOp:
     op = xla::Erf(*inputs[0]);
     break;
+  case IsFiniteOp:
+    op = xla::IsFinite(*inputs[0]);
+    break;
+  case PopulationCountOp:
+    op = xla::PopulationCount(*inputs[0]);
+    break;
 
   // Two-arguments ops
   case AddOp:

diff --git a/cmd/protoc_xla_protos/main.go b/cmd/protoc_xla_protos/main.go
@@ -25,7 +25,7 @@ const (
 )
 
 var protos = []string{
-	"third_party/tsl/tsl/protobuf/dnn.proto",
+	"xla/tsl/protobuf/dnn.proto",
 	"xla/autotune_results.proto",
 	"xla/autotuning.proto",
 	"xla/pjrt/compile_options.proto",

diff --git a/cmd/xlabuilder_codegen/go_ops.go b/cmd/xlabuilder_codegen/go_ops.go
@@ -111,9 +111,16 @@ func TestSimpleOps(t *testing.T) {
 	result, err = And(result, same)
 	require.NoError(t, err, "Failed to build And operation")
 
+	result, err = IsFinite(x)
+	require.NoError(t, err, "Failed to build IsFinite operation")
+	require.Equal(t, dtypes.Bool, result.Shape.DType, "IsFinite should return booleans") 
+
 	require.NoError(t, err, "Failed to build And operation")
 	i, err = Clz(i)
 	require.NoError(t, err, "Failed to build Clz operation")
+
+	i, err = PopulationCount(i)
+	require.NoError(t, err, "Failed to build PopulationCount operation")
 	same, err = Equal(i, i)
 	require.NoError(t, err, "Failed to build Equal operation")
 
@@ -149,7 +156,7 @@ func GenerateSimpleGoOps(opsInfo []OpInfo) {
 	fmt.Printf("Generated %q based on %q\n", fileName, OpTypesFileName)
 
 	// For testing we skip some that require special types.
-	skip := []string{"LogicalNot", "And", "Or", "Xor", "Dot", "Clz", "Real", "Imag", "Conj", "Complex"}
+	skip := []string{"LogicalNot", "And", "Or", "Xor", "Dot", "Clz", "Real", "Imag", "Conj", "Complex", "IsFinite", "PopulationCount"}
 	filteredOps := slices.DeleteFunc(opsInfo, func(info OpInfo) bool {
 		return slices.Index(skip, info.Name) != -1
 	})

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
@@ -1,3 +1,7 @@
+# v0.4.2 -  2024-10-03
+
+* Added `IsFinite` and `PopulationCount` operations.
+
 # v0.4.1 - 2024-09-28
 
 * Added memory layout information in buffer-to-host transfers: required for TPU.

diff --git a/protos/autotuning/autotuning.pb.go b/protos/autotuning/autotuning.pb.go