Add missing tensor data types (unsigned int 16, 32, 64) to PopulateTensorBuffer (#9090)

Author: iwknow

test/cpp/test_tensor.cpp

@@ -62,6 +62,28 @@ TEST_F(TensorTest, TestConversions) {
                                at::TensorOptions(at::kInt));
     EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::Long));
   }
+  {
+    at::Tensor a = at::randint(std::numeric_limits<uint16_t>::min(),
+                               std::numeric_limits<uint16_t>::max(), {2, 2},
+                               at::TensorOptions(at::kUInt16));
+    EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::UInt16));
+    EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::UInt32));
+    EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::UInt64));
+  }
+  {
+    at::Tensor a = at::randint(std::numeric_limits<uint32_t>::min(),
+                               std::numeric_limits<uint32_t>::max(), {2, 2},
+                               at::TensorOptions(at::kUInt32));
+    EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::UInt32));
+    EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::UInt64));
+  }
+  {
+    // The range of uint64_t is too large for randint to generate.
+    at::Tensor a = at::randint(std::numeric_limits<uint32_t>::min(),
+                               std::numeric_limits<uint32_t>::max(), {2, 2},
+                               at::TensorOptions(at::kUInt64));
+    EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::UInt64));
+  }
   {
     at::Tensor a = at::randint(0, 1, {2, 2}, at::TensorOptions(at::kByte));
     EXPECT_TRUE(CheckBidirectionalConversion(a, at::ScalarType::Byte,

torch_xla/csrc/tensor_util.cpp

@@ -701,6 +701,18 @@ void PopulateTensorBuffer(const at::Tensor& tensor,
       TensorToBufferSType<c10::complex<double>>(tensor, dest_shape, dest_buffer,
                                                 dest_buffer_size, device);
       break;
+    case at::ScalarType::UInt16:
+      TensorToBufferSType<uint16_t>(tensor, dest_shape, dest_buffer,
+                                    dest_buffer_size, device);
+      break;
+    case at::ScalarType::UInt32:
+      TensorToBufferSType<uint32_t>(tensor, dest_shape, dest_buffer,
+                                    dest_buffer_size, device);
+      break;
+    case at::ScalarType::UInt64:
+      TensorToBufferSType<uint64_t>(tensor, dest_shape, dest_buffer,
+                                    dest_buffer_size, device);
+      break;
     default:
       XLA_ERROR() << "Tensor type not supported: " << tensor.type();
   }