[XLA] Fix bool support for Array2D/Array3D/Array4D.

The Array?D<T> classes were previously backed by a std::vector<T>, which has an inconsistent bitset implementation for T == bool. This meant that the bool instantiations of Array?D caused compile-time errors. Change the Array classes to be backed by a std::unique_ptr<T[]> instead, since the dimensions of an Array are constant.

PiperOrigin-RevId: 156364397

4 files changed, 115 insertions, 48 deletions

diff --git a/tensorflow/compiler/xla/array2d.h b/tensorflow/compiler/xla/array2d.h
index f885821210..593084a0c1 100644
--- a/tensorflow/compiler/xla/array2d.h
+++ b/tensorflow/compiler/xla/array2d.h
@@ -45,11 +45,15 @@ class Array2D {
 
   // Creates an array of dimensions n1 x n2, uninitialized values.
   Array2D(const int64 n1, const int64 n2)
-      : n1_(n1), n2_(n2), values_(n1 * n2) {}
+      : n1_(n1), n2_(n2), values_(new T[n1 * n2]()) {
+    Fill(T());
+  }
 
   // Creates an array of dimensions n1 x n2, initialized to value.
   Array2D(const int64 n1, const int64 n2, const T value)
-      : n1_(n1), n2_(n2), values_(n1 * n2, value) {}
+      : n1_(n1), n2_(n2), values_(new T[n1 * n2]()) {
+    Fill(value);
+  }
 
   // Creates an array from the given nested initializer list. The outer
   // initializer list is the first dimension; the inner is the second dimension.
@@ -65,16 +69,30 @@ class Array2D {
     }
   }
 
-  T& operator()(const int64 n1, const int64 n2) {
-    CHECK_LT(n1, n1_);
-    CHECK_LT(n2, n2_);
-    return values_[n1 * n2_ + n2];
+  Array2D(const Array2D<T>& other) : Array2D(other.n1(), other.n2()) {
+    std::copy(&other.values_[0], &other.values_[0] + num_elements(),
+              &values_[0]);
+  }
+
+  Array2D<T>& operator=(const Array2D<T>& other) {
+    n1_ = other.n1();
+    n2_ = other.n2();
+    values_.reset(new T[num_elements()]);
+    std::copy(&other.values_[0], &other.values_[0] + num_elements(),
+              &values_[0]);
+    return *this;
+  }
+
+  T& operator()(const int64 i1, const int64 i2) {
+    CHECK_LT(i1, n1_);
+    CHECK_LT(i2, n2_);
+    return values_[i1 * n2_ + i2];
   }
 
-  const T& operator()(const int64 n1, const int64 n2) const {
-    CHECK_LT(n1, n1_);
-    CHECK_LT(n2, n2_);
-    return values_[n1 * n2_ + n2];
+  const T& operator()(const int64 i1, const int64 i2) const {
+    CHECK_LT(i1, n1_);
+    CHECK_LT(i2, n2_);
+    return values_[i1 * n2_ + i2];
   }
 
   // Access to the array's dimensions. height() and width() provide the
@@ -84,15 +102,15 @@ class Array2D {
   int64 n2() const { return n2_; }
   int64 height() const { return n1_; }
   int64 width() const { return n2_; }
-  int64 num_elements() const { return values_.size(); }
+  int64 num_elements() const { return n1_ * n2_; }
 
   // Low-level accessor for stuff like memcmp, handle with care. Returns pointer
   // to the underlying storage of the array (similarly to std::vector::data()).
-  T* data() const { return const_cast<Array2D*>(this)->values_.data(); }
+  T* data() const { return const_cast<Array2D*>(this)->values_.get(); }
 
   // Fills the array with the given value.
   void Fill(const T& value) {
-    std::fill(values_.begin(), values_.end(), value);
+    std::fill(&values_[0], &values_[0] + num_elements(), value);
   }
 
   // Applies f to all cells in this array, in row-major order.
@@ -124,8 +142,8 @@ class Array2D {
     std::mt19937 g(seed);
     std::normal_distribution<double> distribution(mean,
                                                   static_cast<double>(value));
-    for (auto& v : values_) {
-      v = static_cast<T>(distribution(g));
+    for (int64 i = 0; i < num_elements(); ++i) {
+      values_[i] = static_cast<T>(distribution(g));
     }
   }
 
@@ -150,7 +168,7 @@ class Array2D {
  private:
   int64 n1_;
   int64 n2_;
-  std::vector<T> values_;
+  std::unique_ptr<T[]> values_;
 };
 
 // Returns a linspace-populated Array2D in the range [from, to] (inclusive)
diff --git a/tensorflow/compiler/xla/array2d_test.cc b/tensorflow/compiler/xla/array2d_test.cc
index 0d502eaf3b..795d50ca5b 100644
--- a/tensorflow/compiler/xla/array2d_test.cc
+++ b/tensorflow/compiler/xla/array2d_test.cc
@@ -84,6 +84,17 @@ TEST(Array2dTest, IndexingReadWrite) {
   EXPECT_EQ(arr(1, 2), 61);
 }
 
+TEST(Array2dTest, IndexingReadWriteBool) {
+  Array2D<bool> arr = {{false, true, false}, {true, true, false}};
+
+  EXPECT_EQ(arr(1, 1), true);
+  EXPECT_EQ(arr(1, 2), false);
+  arr(1, 1) = false;
+  arr(1, 2) = true;
+  EXPECT_EQ(arr(1, 1), false);
+  EXPECT_EQ(arr(1, 2), true);
+}
+
 TEST(Array2dTest, Fill) {
   Array2D<int> fullof7(2, 3, 7);
   for (int64 n1 = 0; n1 < fullof7.n1(); ++n1) {
diff --git a/tensorflow/compiler/xla/array3d.h b/tensorflow/compiler/xla/array3d.h
index 654af8f030..124ccd1975 100644
--- a/tensorflow/compiler/xla/array3d.h
+++ b/tensorflow/compiler/xla/array3d.h
@@ -20,9 +20,9 @@ limitations under the License.
 #include <functional>
 #include <initializer_list>
 #include <iterator>
+#include <memory>
 #include <numeric>
 #include <random>
-#include <vector>
 
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/core/platform/logging.h"
@@ -39,11 +39,15 @@ class Array3D {
  public:
   // Creates an array of dimensions n1 x n2 x n3, uninitialized values.
   Array3D(const int64 n1, const int64 n2, const int64 n3)
-      : n1_(n1), n2_(n2), n3_(n3), values_(n1 * n2 * n3) {}
+      : n1_(n1), n2_(n2), n3_(n3), values_(new T[n1 * n2 * n3]) {
+    Fill(T());
+  }
 
   // Creates an array of dimensions n1 x n2 x n3, initialized to value.
   Array3D(const int64 n1, const int64 n2, const int64 n3, const T value)
-      : n1_(n1), n2_(n2), n3_(n3), values_(n1 * n2 * n3, value) {}
+      : n1_(n1), n2_(n2), n3_(n3), values_(new T[n1 * n2 * n3]) {
+    Fill(value);
+  }
 
   // Creates an array from the given nested initializer list. The outer
   // initializer list is the first dimension, and so on.
@@ -69,34 +73,50 @@ class Array3D {
     }
   }
 
-  T& operator()(const int64 n1, const int64 n2, const int64 n3) {
-    CHECK_LT(n1, n1_);
-    CHECK_LT(n2, n2_);
-    CHECK_LT(n3, n3_);
-    return values_[n1 * n2_ * n3_ + n2 * n3_ + n3];
+  Array3D(const Array3D<T>& other)
+      : Array3D(other.n1(), other.n2(), other.n3()) {
+    std::copy(&other.values_[0], &other.values_[0] + num_elements(),
+              &values_[0]);
+  }
+
+  Array3D<T>& operator=(const Array3D<T>& other) {
+    n1_ = other.n1();
+    n2_ = other.n2();
+    n3_ = other.n3();
+    values_.reset(new T[num_elements()]);
+    std::copy(&other.values_[0], &other.values_[0] + num_elements(),
+              &values_[0]);
+    return *this;
+  }
+
+  T& operator()(const int64 i1, const int64 i2, const int64 i3) {
+    CHECK_LT(i1, n1_);
+    CHECK_LT(i2, n2_);
+    CHECK_LT(i3, n3_);
+    return values_[i1 * n2_ * n3_ + i2 * n3_ + i3];
   }
 
-  const T& operator()(const int64 n1, const int64 n2, const int64 n3) const {
-    CHECK_LT(n1, n1_);
-    CHECK_LT(n2, n2_);
-    CHECK_LT(n3, n3_);
-    return values_[n1 * n2_ * n3_ + n2 * n3_ + n3];
+  const T& operator()(const int64 i1, const int64 i2, const int64 i3) const {
+    CHECK_LT(i1, n1_);
+    CHECK_LT(i2, n2_);
+    CHECK_LT(i3, n3_);
+    return values_[i1 * n2_ * n3_ + i2 * n3_ + i3];
   }
 
   // Access to the array's dimensions.
   int64 n1() const { return n1_; }
   int64 n2() const { return n2_; }
   int64 n3() const { return n3_; }
-  int64 num_elements() const { return values_.size(); }
+  int64 num_elements() const { return n1_ * n2_ * n3_; }
 
   // Fills the array with the given value.
   void Fill(const T& value) {
-    std::fill(values_.begin(), values_.end(), value);
+    std::fill(&values_[0], &values_[0] + num_elements(), value);
   }
 
   // Fills the array with sequentially increasing values.
   void FillIota(const T& value) {
-    std::iota(values_.begin(), values_.end(), value);
+    std::iota(&values_[0], &values_[0] + num_elements(), value);
   }
 
   // Fills the array with random normal values with a mean of 0 and standard
@@ -106,8 +126,8 @@ class Array3D {
     std::mt19937 g(seed);
     std::normal_distribution<double> distribution(mean,
                                                   static_cast<double>(value));
-    for (auto& v : values_) {
-      v = static_cast<T>(distribution(g));
+    for (int64 i = 0; i < num_elements(); ++i) {
+      values_[i] = static_cast<T>(distribution(g));
     }
   }
 
@@ -115,7 +135,7 @@ class Array3D {
   int64 n1_;
   int64 n2_;
   int64 n3_;
-  std::vector<T> values_;
+  std::unique_ptr<T[]> values_;
 };
 
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/array4d.h b/tensorflow/compiler/xla/array4d.h
index 199ad2baae..1c6ba1f519 100644
--- a/tensorflow/compiler/xla/array4d.h
+++ b/tensorflow/compiler/xla/array4d.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include <functional>
 #include <initializer_list>
 #include <iterator>
+#include <memory>
 #include <numeric>
 #include <random>
 #include <string>
@@ -60,15 +61,15 @@ class Array4D {
         depth_(depth),
         height_(height),
         width_(width),
-        values_(planes * depth * height * width) {}
+        values_(new T[planes * depth * height * width]) {
+    Fill(T());
+  }
 
   // Creates a 4D array, initalized to value.
   Array4D(int64 planes, int64 depth, int64 height, int64 width, T value)
-      : planes_(planes),
-        depth_(depth),
-        height_(height),
-        width_(width),
-        values_(planes * depth * height * width, value) {}
+      : Array4D(planes, depth, height, width) {
+    Fill(value);
+  }
 
   // Creates a 4D array, filled with values.
   //
@@ -111,6 +112,23 @@ class Array4D {
     }
   }
 
+  Array4D(const Array4D<T>& other)
+      : Array4D(other.planes(), other.depth(), other.height(), other.width()) {
+    std::copy(&other.values_[0], &other.values_[0] + num_elements(),
+              &values_[0]);
+  }
+
+  Array4D<T>& operator=(const Array4D<T>& other) {
+    planes_ = other.planes();
+    depth_ = other.depth();
+    height_ = other.height();
+    width_ = other.width();
+    values_.reset(new T[num_elements()]);
+    std::copy(&other.values_[0], &other.values_[0] + num_elements(),
+              &values_[0]);
+    return *this;
+  }
+
   T& operator()(int64 plane, int64 depth, int64 height, int64 width) {
     CHECK_LT(plane, planes_);
     CHECK_LT(depth, depth_);
@@ -135,24 +153,24 @@ class Array4D {
   int64 n3() const { return height_; }
   int64 n2() const { return depth_; }
   int64 n1() const { return planes_; }
-  int64 num_elements() const { return values_.size(); }
+  int64 num_elements() const { return width_ * height_ * depth_ * planes_; }
 
   // Sets all the values in the array to values.
   template <typename Container = std::initializer_list<T>>
   void SetValues(const Container& container) {
     CHECK_EQ(std::distance(std::begin(container), std::end(container)),
              num_elements());
-    values_.assign(std::begin(container), std::end(container));
+    std::copy(std::begin(container), std::end(container), &values_[0]);
   }
 
   // Fills the array with the given value.
   void Fill(const T& value) {
-    std::fill(values_.begin(), values_.end(), value);
+    std::fill(&values_[0], &values_[0] + num_elements(), value);
   }
 
   // Fills the array with iota.
   void FillIota(const T& value) {
-    std::iota(values_.begin(), values_.end(), value);
+    std::iota(&values_[0], &values_[0] + num_elements(), value);
   }
 
   // Fills the array with random variable with a deviation of value and a mean
@@ -162,8 +180,8 @@ class Array4D {
     std::mt19937 g(seed);
     std::normal_distribution<double> distribution(mean,
                                                   static_cast<double>(value));
-    for (auto& v : values_) {
-      v = static_cast<T>(distribution(g));
+    for (int64 i = 0; i < num_elements(); ++i) {
+      values_[i] = static_cast<T>(distribution(g));
     }
   }
 
@@ -268,7 +286,7 @@ class Array4D {
   int64 depth_;
   int64 height_;
   int64 width_;
-  std::vector<T> values_;
+  std::unique_ptr<T[]> values_;
 };
 
 }  // namespace xla