| /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
| |
| 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 "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" |
| #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" |
| #include "tensorflow/compiler/xla/service/hlo_parser.h" |
| |
| namespace xla { |
| namespace cpu { |
| namespace { |
| class CpuKeyValueSortTest : public CpuCodegenTest {}; |
| |
| TEST_F(CpuKeyValueSortTest, SortR1) { |
| const string hlo_text = R"( |
| HloModule KeyValueSort |
| |
| ENTRY main { |
| a = f32[10] parameter(0) |
| |
| ROOT result = f32[10] sort(f32[10] a), dimensions={0} |
| } |
| )"; |
| |
| string filecheck_pattern = R"( |
| CHECK: call void @__xla_cpu_runtime_KeyValueSort |
| )"; |
| |
| TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module, |
| ParseHloString(hlo_text)); |
| |
| CpuAotCompilationOptions options{ |
| /*triple=*/"x86_64", /*cpu_name=*/"", /*features=*/"", |
| /*entry_point_name=*/"entry", |
| /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; |
| |
| CompileAheadOfTimeAndVerifyIr(std::move(module), options, filecheck_pattern, |
| /*match_optimized_ir=*/true); |
| } |
| |
| } // namespace |
| } // namespace cpu |
| } // namespace xla |