Improve S1/U1 support

PiperOrigin-RevId: 708444896
makesoftwaresafe · Dec 20, 2024 · 950d61d · 950d61d
1 parent b198703
commit 950d61d
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 0 deletions.
diff --git a/third_party/xla/xla/hlo/evaluator/BUILD b/third_party/xla/xla/hlo/evaluator/BUILD
@@ -30,6 +30,7 @@ cc_library(
         "hlo_evaluator_typed_visitor_float.cc",
         "hlo_evaluator_typed_visitor_float8.cc",
         "hlo_evaluator_typed_visitor_half.cc",
+        "hlo_evaluator_typed_visitor_int1.cc",
         "hlo_evaluator_typed_visitor_int16.cc",
         "hlo_evaluator_typed_visitor_int2.cc",
         "hlo_evaluator_typed_visitor_int32.cc",

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
@@ -1716,12 +1716,14 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
 // instantiating it.  We explicitly instantiate this class in the various
 // hlo_evaluator_typed_visitor*.cc files.
 extern template class HloEvaluatorTypedVisitor<bool>;
+extern template class HloEvaluatorTypedVisitor<u1, uint64_t>;
 extern template class HloEvaluatorTypedVisitor<u2, uint64_t>;
 extern template class HloEvaluatorTypedVisitor<u4, uint64_t>;
 extern template class HloEvaluatorTypedVisitor<uint8_t, uint64_t>;
 extern template class HloEvaluatorTypedVisitor<uint16_t, uint64_t>;
 extern template class HloEvaluatorTypedVisitor<uint32_t, uint64_t>;
 extern template class HloEvaluatorTypedVisitor<uint64_t>;
+extern template class HloEvaluatorTypedVisitor<s1, int64_t>;
 extern template class HloEvaluatorTypedVisitor<s2, int64_t>;
 extern template class HloEvaluatorTypedVisitor<s4, int64_t>;
 extern template class HloEvaluatorTypedVisitor<int8_t, int64_t>;

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor_int1.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor_int1.cc
@@ -0,0 +1,25 @@
+/* Copyright 2018 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "xla/hlo/evaluator/hlo_evaluator.h"
+#include "xla/hlo/evaluator/hlo_evaluator_typed_visitor.h"
+#include "xla/types.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<s1, int64_t>;
+template class HloEvaluatorTypedVisitor<u1, uint64_t>;
+}  // namespace xla
diff --git a/third_party/xla/xla/literal_util.cc b/third_party/xla/xla/literal_util.cc
@@ -472,6 +472,11 @@ void PopulateWithRandomIntegralDataWithBounds(Literal* literal,
   return ConvertType<int32_t, float>(s32_literal);
 }
 
+/* static */ Literal LiteralUtil::ConvertS32ToS1(
+    const LiteralSlice& s32_literal) {
+  return ConvertType<int32_t, tsl::int1>(s32_literal);
+}
+
 /* static */ Literal LiteralUtil::CreateToken() {
   return Literal(ShapeUtil::MakeTokenShape());
 }

diff --git a/third_party/xla/xla/literal_util.h b/third_party/xla/xla/literal_util.h
@@ -252,6 +252,7 @@ class LiteralUtil {
   static Literal ConvertF64ToBF16(const LiteralSlice& f64_literal);
   static Literal ConvertF64ToF32(const LiteralSlice& f64_literal);
   static Literal ConvertS32ToF32(const LiteralSlice& s32_literal);
+  static Literal ConvertS32ToS1(const LiteralSlice& s32_literal);
 
   // Creates a scalar literal whose value is the maximum value of a given
   // literal slice.