1effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 3effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasLicensed under the Apache License, Version 2.0 (the "License"); 4effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dasyou may not use this file except in compliance with the License. 5effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasYou may obtain a copy of the License at 6effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 7effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das http://www.apache.org/licenses/LICENSE-2.0 8effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 9effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasUnless required by applicable law or agreed to in writing, software 10effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dasdistributed under the License is distributed on an "AS IS" BASIS, 11effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasSee the License for the specific language governing permissions and 13effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Daslimitations under the License. 14effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das==============================================================================*/ 15effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 16effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das#include "tensorflow/compiler/xla/service/cpu/external_constant_pool.h" 17effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das#include "tensorflow/compiler/xla/layout_util.h" 18effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das#include "tensorflow/core/platform/test.h" 19effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 20effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dasnamespace xla { 21effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dasnamespace cpu { 22effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dasnamespace { 23effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dasclass ExternalConstantPoolTest : public ::testing::Test {}; 24effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 25effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Dastemplate <typename T> 26effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasT GetFromBuffer(const uint8* buffer, int64 index) { 27effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das T result; 28effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das std::memcpy(&result, buffer + index * sizeof(T), sizeof(T)); 29effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das return result; 30effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} 31effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 32effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasTEST(ExternalConstantPoolTest, Basic) { 33effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das ExternalConstantPool constant_pool; 34effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(constant_pool.Find("name-0"), nullptr); 35effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das const auto literal = Literal::CreateR2({{1, 2}, {3, 4}}); 36effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das constant_pool.Insert("name-0", *literal, 4); 37effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das const uint8* constant = constant_pool.Find("name-0"); 38effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das ASSERT_NE(constant, nullptr); 39effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 40effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 0), 1); 41effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 1), 2); 42effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 2), 3); 43effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 3), 4); 44effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 45effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(constant_pool.Find("name-1"), nullptr); 46effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} 47effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 48effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasTEST(ExternalConstantPoolTest, RowMinorLayout) { 49effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das ExternalConstantPool constant_pool; 50effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(constant_pool.Find("name-0"), nullptr); 51effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das const auto literal = Literal::CreateR2WithLayout( 52effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das {{1, 2}, {3, 4}}, LayoutUtil::MakeLayout({0, 1})); 53effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das constant_pool.Insert("name-0", *literal, 4); 54effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das const uint8* constant = constant_pool.Find("name-0"); 55effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das ASSERT_NE(constant, nullptr); 56effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 57effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 0), 1); 58effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 1), 3); 59effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 2), 2); 60effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(GetFromBuffer<int32>(constant, 3), 4); 61effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} 62effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 63effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy DasTEST(ExternalConstantPoolTest, Alignment) { 64effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das ExternalConstantPool constant_pool; 65effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(constant_pool.Find("name-0"), nullptr); 66effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 67effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das for (int i = 0; i < 8; i++) { 68effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das int64 alignment = 1 << i; 69effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das string name = tensorflow::strings::StrCat("name-", i); 70effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 71effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das const auto literal = Literal::CreateR2({{1, 2}, {3, 4}}); 72effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das constant_pool.Insert(name, *literal, alignment); 73effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 74effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das const uint8* constant = constant_pool.Find(name); 75effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das ASSERT_NE(constant, nullptr); 76effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das EXPECT_EQ(reinterpret_cast<intptr_t>(constant) % alignment, 0); 77effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das } 78effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} 79effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das 80effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} // namespace 81effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} // namespace cpu 82effb22e8a44763901ee2cf55c30290f0b1edb570Sanjoy Das} // namespace xla 83