chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,18 @@
|
||||
if(NOT WITH_CPP_TEST)
|
||||
return()
|
||||
endif()
|
||||
add_subdirectory(auto_parallel)
|
||||
add_subdirectory(phi)
|
||||
add_subdirectory(jit)
|
||||
add_subdirectory(new_executor)
|
||||
add_subdirectory(prim)
|
||||
add_subdirectory(imperative)
|
||||
add_subdirectory(pir)
|
||||
add_subdirectory(inference)
|
||||
add_subdirectory(eager)
|
||||
add_subdirectory(fluid)
|
||||
add_subdirectory(utils)
|
||||
add_subdirectory(compat)
|
||||
if(WITH_CINN)
|
||||
add_subdirectory(cinn)
|
||||
endif()
|
||||
@@ -0,0 +1,96 @@
|
||||
if(WIN32)
|
||||
cc_test(
|
||||
device_mesh_test
|
||||
SRCS device_mesh_test.cc
|
||||
DEPS type_info)
|
||||
cc_test(
|
||||
process_mesh_test
|
||||
SRCS process_mesh_test.cc
|
||||
DEPS type_info)
|
||||
else()
|
||||
cc_test(device_mesh_test SRCS device_mesh_test.cc)
|
||||
cc_test(process_mesh_test SRCS process_mesh_test.cc)
|
||||
endif()
|
||||
|
||||
cc_test(
|
||||
dist_attr_test
|
||||
SRCS dist_attr_test.cc
|
||||
DEPS proto_desc)
|
||||
|
||||
if(WITH_DISTRIBUTE)
|
||||
cc_library(
|
||||
spmd_rule_test_util
|
||||
SRCS spmd_rule_test_util.cc
|
||||
DEPS gtest)
|
||||
cc_test(
|
||||
dist_tensor_test
|
||||
SRCS dist_tensor_test.cc
|
||||
DEPS phi common)
|
||||
|
||||
paddle_test(spmd_rule_test SRCS spmd_rule_test.cc DEPS spmd_rule_test_util
|
||||
phi)
|
||||
|
||||
paddle_test(softmax_grad_spmd_rule_test SRCS softmax_grad_spmd_rule_test.cc
|
||||
DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(tile_spmd_rule_test SRCS tile_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
paddle_test(tile_co_shard_spmd_rule_test SRCS tile_co_shard_spmd_rule_test.cc
|
||||
DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(
|
||||
fused_linear_param_grad_add_spmd_rule_test SRCS
|
||||
fused_linear_param_grad_add_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(
|
||||
cross_entropy_softmax_spmd_rule_test SRCS
|
||||
cross_entropy_softmax_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(expand_spmd_rule_test SRCS expand_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(expand_as_spmd_rule_test SRCS expand_as_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(matmul_co_shard_spmd_rule_test SRCS
|
||||
matmul_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(custom_op_spmd_rule_test SRCS custom_op_spmd_rule_test.cc DEPS
|
||||
spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(fused_rms_norm_spmd_rule_test SRCS
|
||||
fused_rms_norm_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
paddle_test(moe_gate_dispatch_spmd_rule_test SRCS
|
||||
moe_gate_dispatch_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(moe_combine_spmd_rule_test SRCS moe_combine_spmd_rule_test.cc
|
||||
DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(softmax_co_shard_spmd_rule_test SRCS
|
||||
softmax_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(
|
||||
index_select_co_shard_spmd_rule_test SRCS
|
||||
index_select_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(reshape_co_shard_spmd_rule_test SRCS
|
||||
reshape_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
paddle_test(argsort_co_shard_spmd_rule_test SRCS
|
||||
argsort_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
paddle_test(transpose_co_shard_spmd_rule_test SRCS
|
||||
transpose_co_shard_spmd_rule_test.cc DEPS spmd_rule_test_util phi)
|
||||
|
||||
endif()
|
||||
|
||||
if(WIN32)
|
||||
cc_test(
|
||||
dist_mapper_test
|
||||
SRCS dist_mapper_test.cc
|
||||
DEPS type_info)
|
||||
else()
|
||||
cc_test(
|
||||
dist_mapper_test
|
||||
SRCS dist_mapper_test.cc
|
||||
DEPS phi)
|
||||
endif()
|
||||
@@ -0,0 +1,205 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct ArgSortTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_indices_dims_mapping;
|
||||
|
||||
// unused attribute
|
||||
bool descending = true;
|
||||
bool stable = true;
|
||||
};
|
||||
|
||||
struct ArgSortGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> indices_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_indices_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
// unused attribute
|
||||
bool descending = true;
|
||||
bool stable = true;
|
||||
};
|
||||
|
||||
TEST(ArgSortInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<ArgSortTestCase> test_cases = {
|
||||
// shape = [16, 32, 48], axis = -1
|
||||
// [[0,1],[2],[]] -> [[],[2],[]], [[],[2],[]]
|
||||
{{16, 32, 48},
|
||||
{{0, 1}, {2}, {}},
|
||||
-1,
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}}},
|
||||
|
||||
// shape = [16, 32, 48], axis = 2
|
||||
// [[0],[],[1,2]] -> [[0],[],[]], [[0],[],[]]
|
||||
{{16, 32, 48},
|
||||
{{0}, {}, {1, 2}},
|
||||
2,
|
||||
{{0}, {}, {}},
|
||||
{{0}, {}, {}},
|
||||
{{0}, {}, {}}},
|
||||
|
||||
// shape = [10, 32, 48, 24], axis = 1
|
||||
// [[0,1],[2],[],[]] -> [[0,1],[],[],[]], [[0,1],[],[],[]]
|
||||
{{10, 32, 48, 24},
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{0, 1}, {}, {}, {}}}};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::ArgSortInferSpmd(
|
||||
x, tc.axis, tc.descending, tc.stable);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(2));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[1],
|
||||
tc.expected_indices_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(ArgSortGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<ArgSortGradTestCase> test_cases = {
|
||||
// shape = [16, 32, 48], axis = -1
|
||||
// [[0,1],[2],[]], [[0,1],[2],[]], [[0,1],[2],[]] -> [[0,1],[2],[]],
|
||||
// [[0,1],[2],[]], [[0,1],[2],[]], [[0,1],[2],[]]
|
||||
{{16, 32, 48},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
-1,
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {2}, {}}},
|
||||
// axis = 2
|
||||
// [[0,1],[],[2]], [[0,1],[],[2]], [[0,1],[],[2]] -> [[0,1],[],[]],
|
||||
// [[0,1],[],[]], [[0,1],[],[]], [[0,1],[],[]]
|
||||
{{16, 32, 48},
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
2,
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
// [10, 32, 48, 24], axis = 1
|
||||
// [[0],[1,2],[]], [[0],[1,2],[]], [[0],[1,2],[]] -> [[0],[],[]],
|
||||
// [[0],[],[]], [[0],[],[]], [[0],[],[]]
|
||||
{{10, 32, 48, 24},
|
||||
{{0}, {1, 2}, {}, {}},
|
||||
{{0}, {1, 2}, {}, {}},
|
||||
{{0}, {1, 2}, {}, {}},
|
||||
1,
|
||||
{{0}, {}, {}, {}},
|
||||
{{0}, {}, {}, {}},
|
||||
{{0}, {}, {}, {}},
|
||||
{{0}, {}, {}, {}}}};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr indices_dist_attr = TensorDistAttr();
|
||||
indices_dist_attr.set_process_mesh(process_mesh);
|
||||
indices_dist_attr.set_dims_mapping(tc.indices_dims_mapping);
|
||||
indices_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor indices = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), indices_dist_attr);
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), x_dist_attr);
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.input_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::ArgSortGradInferSpmd(
|
||||
indices, x, out_grad, tc.axis, tc.descending, tc.stable);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_indices_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[2],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,127 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(CrossEntropyInferSpmd, Ctor) {
|
||||
std::vector<int64_t> x_shape = {32, 48};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
TensorDistAttr label_dist_attr = TensorDistAttr();
|
||||
label_dist_attr.set_process_mesh(process_mesh);
|
||||
label_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
label_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
// forward
|
||||
{
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor label(phi::make_ddim(x_shape),
|
||||
label_dist_attr);
|
||||
int axis = 1;
|
||||
|
||||
auto spmdinfo =
|
||||
CrossEntropyWithSoftmaxInferSpmd(x, label, false, true, true, 1, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 2UL);
|
||||
check_dim_mapping(spmdinfo.first[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[1], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[1], {0, -1});
|
||||
check_partial_dims(spmdinfo.second[0], {});
|
||||
|
||||
VLOG(4) << "Test CrossEntropyWithSoftmaxInferSpmd sharding on other axes."
|
||||
<< std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
// test sharding along softmax axis.
|
||||
{
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
|
||||
label_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor label(phi::make_ddim(x_shape),
|
||||
label_dist_attr);
|
||||
int axis = 1;
|
||||
|
||||
auto spmdinfo =
|
||||
CrossEntropyWithSoftmaxInferSpmd(x, label, false, true, true, 1, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 2UL);
|
||||
check_dim_mapping(spmdinfo.first[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[1], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[1], {0, -1});
|
||||
check_partial_dims(spmdinfo.second[0], {});
|
||||
|
||||
VLOG(4) << "Test CrossEntropyWithSoftmaxInferSpmd sharding on other axes."
|
||||
<< std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
// backward
|
||||
{
|
||||
std::vector<int64_t> loss_shape = {32, 1};
|
||||
// Sharding along softmax axis.
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>{0, 1});
|
||||
label_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
|
||||
auto label = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
label_dist_attr);
|
||||
auto softmax =
|
||||
phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
|
||||
auto loss_dist_attr = x_dist_attr;
|
||||
loss_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
auto loss_grad = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(loss_shape), x_dist_attr);
|
||||
|
||||
int axis = 1;
|
||||
auto spmdinfo = CrossEntropyWithSoftmaxGradInferSpmd(
|
||||
label, softmax, loss_grad, true, true, true, 1, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 3UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
check_dim_mapping(spmdinfo.first[0], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[1], {0, -1});
|
||||
check_dim_mapping(spmdinfo.first[2], {0, -1});
|
||||
check_dim_mapping(spmdinfo.second[0], {0, -1});
|
||||
check_partial_dims(spmdinfo.second[0], {});
|
||||
|
||||
VLOG(4)
|
||||
<< "Test CrossEntropyWithSoftmaxGradInferSpmd sharding on softmax axis."
|
||||
<< std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,89 @@
|
||||
/* Copyright (c) 2023 PaddlePaddle 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 "paddle/phi/api/ext/op_meta_info.h"
|
||||
#include "paddle/phi/api/ext/spmd_infer.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
TEST(CustomOp, Ctor) {
|
||||
// test with concat rule
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<std::vector<int64_t>> shapes = {
|
||||
{16, 16, 16}, {4, 16, 16}, {2, 16, 16}};
|
||||
std::vector<std::vector<int64_t>> dim_mappings = {
|
||||
{-1, 0, 1}, {-1, 1, 0}, {-1, -1, 0}};
|
||||
std::vector<std::vector<int64_t>> partial_status = {{}, {}, {1}};
|
||||
|
||||
auto build_inputs = [&] {
|
||||
std::vector<phi::distributed::DistMetaTensor> inputs;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(dim_mappings[i]);
|
||||
t_dist_attr.set_dynamic_dims({false, false, false});
|
||||
auto input = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(shapes[i]), t_dist_attr);
|
||||
inputs.push_back(input);
|
||||
}
|
||||
return inputs;
|
||||
};
|
||||
|
||||
// test 1, inputs are aligned according to cost, and partial status is cleared
|
||||
auto inputs = build_inputs();
|
||||
|
||||
auto forward_spmd_func =
|
||||
PD_INFER_SPMD_RULE(phi::distributed::ConcatInferSpmd);
|
||||
int axis = 0;
|
||||
std::vector<CustomSpmdInferTensorArg> infer_inputs = {inputs};
|
||||
std::vector<CustomSpmdInferAttrArg> attrs = {axis};
|
||||
|
||||
auto inferred_dist_attrs = forward_spmd_func(infer_inputs, attrs);
|
||||
// list of tensor => single tensor
|
||||
EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
|
||||
inferred_dist_attrs.first[0]));
|
||||
EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
|
||||
inferred_dist_attrs.second[0]));
|
||||
auto& inputs_infer1 =
|
||||
PADDLE_GET_CONST(std::vector<phi::distributed::TensorDistAttr>,
|
||||
inferred_dist_attrs.first[0]);
|
||||
|
||||
for (auto e : inputs_infer1) {
|
||||
check_dim_mapping(e, {-1, 1, 0});
|
||||
check_partial_dims(e, {});
|
||||
}
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {-1, 1, 0});
|
||||
check_partial_dims(inferred_dist_attrs.second[0], {});
|
||||
}
|
||||
|
||||
TEST(CustomOp, Register) {
|
||||
OpMetaInfoBuilder builder("test_custom_op_spmd", 0);
|
||||
auto iter = OpMetaInfoMap::Instance().GetMap().find("test_custom_op_spmd");
|
||||
EXPECT_TRUE(iter != OpMetaInfoMap::Instance().GetMap().end());
|
||||
EXPECT_TRUE(OpMetaInfoHelper::GetInferSpmdFn(iter->second[0]) == nullptr);
|
||||
builder.SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::ConcatInferSpmd));
|
||||
EXPECT_TRUE(OpMetaInfoHelper::GetInferSpmdFn(iter->second[0]) != nullptr);
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,96 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 <iostream>
|
||||
#include <sstream>
|
||||
|
||||
#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(DeviceMesh, Ctor) {
|
||||
std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int64_t> device_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
std::string device_type = "GPU";
|
||||
int64_t size = shape[0] * shape[1];
|
||||
|
||||
DeviceMesh device_mesh("mesh", shape, device_ids, dim_names);
|
||||
for (int64_t i = 0; i < shape[0]; ++i) {
|
||||
for (int64_t j = 0; j < shape[1]; ++j) {
|
||||
int64_t global_id = i * shape[1] + j;
|
||||
int64_t local_id = j;
|
||||
int64_t machine_id = i;
|
||||
device_mesh.add_device(
|
||||
Device(global_id, local_id, machine_id, device_type));
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < size; ++i) {
|
||||
for (int64_t j = 0; j < size; ++j) {
|
||||
device_mesh.add_link(Link(i, j, "NVL"));
|
||||
}
|
||||
}
|
||||
|
||||
EXPECT_EQ(device_mesh.name(), "mesh");
|
||||
EXPECT_EQ(device_mesh.shape(), shape);
|
||||
EXPECT_EQ(device_mesh.device_ids(), device_ids);
|
||||
EXPECT_EQ(device_mesh.dim_names()[0], "x");
|
||||
EXPECT_EQ(device_mesh.dim_names()[1], "y");
|
||||
EXPECT_EQ(device_mesh.device_type(), device_type);
|
||||
EXPECT_EQ(device_mesh.size(), size);
|
||||
EXPECT_EQ(device_mesh.ndim(), static_cast<int64_t>(shape.size()));
|
||||
EXPECT_EQ(device_mesh.dim_size(0), shape[0]);
|
||||
EXPECT_EQ(device_mesh.dim_size(-1), shape[1]);
|
||||
EXPECT_EQ(device_mesh.dim_size("x"), shape[0]);
|
||||
EXPECT_EQ(device_mesh.dim_size("y"), shape[1]);
|
||||
EXPECT_EQ(device_mesh.empty(), false);
|
||||
EXPECT_EQ(device_mesh.contains(0), true);
|
||||
EXPECT_EQ(device_mesh.contains(6), false);
|
||||
EXPECT_EQ(device_mesh.device(3).global_id(), 3);
|
||||
EXPECT_EQ(device_mesh.device(3).local_id(), 0);
|
||||
EXPECT_EQ(device_mesh.device(3).machine_id(), 1);
|
||||
EXPECT_EQ(device_mesh.device(3).type(), "GPU");
|
||||
EXPECT_EQ(device_mesh.link(3, 4).source_id(), 3);
|
||||
EXPECT_EQ(device_mesh.link(3, 4).target_id(), 4);
|
||||
EXPECT_EQ(device_mesh.link(3, 4).type(), "NVL");
|
||||
for (int64_t i = 0; i < shape[0]; ++i) {
|
||||
for (int64_t j = 0; j < shape[1]; ++j) {
|
||||
int64_t global_id = i * shape[1] + j;
|
||||
int64_t local_id = j;
|
||||
int64_t machine_id = i;
|
||||
auto device = device_mesh.devices().at(global_id);
|
||||
EXPECT_EQ(device, Device(global_id, local_id, machine_id, device_type));
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < size; ++i) {
|
||||
for (int64_t j = 0; j < size; ++j) {
|
||||
EXPECT_EQ(device_mesh.links().at(i).at(j), Link(i, j, "NVL"));
|
||||
}
|
||||
}
|
||||
std::stringstream sstream;
|
||||
sstream << device_mesh;
|
||||
EXPECT_EQ(sstream.str(), device_mesh.to_string());
|
||||
auto proto = phi::distributed::to_proto(device_mesh);
|
||||
DeviceMesh new_device_mesh = DeviceMesh::from_proto(proto);
|
||||
EXPECT_EQ(device_mesh, new_device_mesh);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,182 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 <iostream>
|
||||
#include <sstream>
|
||||
#include "glog/logging.h"
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/fluid/distributed/auto_parallel/dist_attr.h"
|
||||
#include "paddle/fluid/framework/block_desc.h"
|
||||
#include "paddle/fluid/framework/op_desc.h"
|
||||
#include "paddle/fluid/framework/program_desc.h"
|
||||
#include "paddle/fluid/framework/var_desc.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using paddle::framework::ProgramDesc;
|
||||
using paddle::framework::VarDesc;
|
||||
|
||||
using paddle::distributed::auto_parallel::get_tensor_shape;
|
||||
using paddle::distributed::auto_parallel::OperatorDistAttr;
|
||||
|
||||
TEST(DistAttr, ctor) {
|
||||
ProgramDesc program;
|
||||
auto* global_block = program.MutableBlock(0);
|
||||
auto* x = global_block->Var("X");
|
||||
x->SetType(paddle::framework::proto::VarType::DENSE_TENSOR);
|
||||
x->SetLoDLevel(0);
|
||||
x->SetDataType(paddle::framework::proto::VarType::FP32);
|
||||
x->SetShape({1000, 784});
|
||||
|
||||
auto* y = global_block->Var("Y");
|
||||
y->SetType(paddle::framework::proto::VarType::DENSE_TENSOR);
|
||||
y->SetLoDLevel(0);
|
||||
y->SetDataType(paddle::framework::proto::VarType::FP32);
|
||||
y->SetShape({784, 100});
|
||||
|
||||
auto* op = global_block->AppendOp();
|
||||
op->SetType("mul");
|
||||
op->SetInput("X", {x->Name()});
|
||||
op->SetInput("Y", {y->Name()});
|
||||
|
||||
auto* out = global_block->Var("Out");
|
||||
out->SetType(paddle::framework::proto::VarType::DENSE_TENSOR);
|
||||
out->SetShape({1000, 100});
|
||||
op->SetOutput("Out", {out->Name()});
|
||||
|
||||
auto get_dist_attr = [](const VarDesc* var_desc) {
|
||||
auto shape = get_tensor_shape(var_desc);
|
||||
return TensorDistAttr(shape);
|
||||
};
|
||||
|
||||
std::vector<int64_t> shape = {2, 4};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> shape2 = {2, 2};
|
||||
std::vector<int64_t> process_ids2 = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names2 = {"a", "b"};
|
||||
ProcessMesh process_mesh2(shape2, process_ids2, dim_names2);
|
||||
|
||||
auto x_dist_attr = get_dist_attr(x);
|
||||
auto y_dist_attr = get_dist_attr(y);
|
||||
auto out_dist_attr = get_dist_attr(out);
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
x_dist_attr.set_batch_dim(0);
|
||||
x_dist_attr.set_chunk_id(0);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({true, false}));
|
||||
x_dist_attr.mark_annotated("process_mesh");
|
||||
x_dist_attr.mark_annotated("dims_mapping");
|
||||
EXPECT_EQ(x_dist_attr.process_mesh(), process_mesh);
|
||||
EXPECT_EQ(x_dist_attr.dims_mapping(), std::vector<int64_t>({0, -1}));
|
||||
EXPECT_EQ(x_dist_attr.batch_dim(), 0);
|
||||
EXPECT_EQ(x_dist_attr.chunk_id(), 0);
|
||||
EXPECT_EQ(x_dist_attr.dynamic_dims(), std::vector<bool>({true, false}));
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("process_mesh"), true);
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("dims_mapping"), true);
|
||||
EXPECT_EQ(x_dist_attr.verify(get_tensor_shape(x)), true);
|
||||
x_dist_attr.clear_annotated();
|
||||
EXPECT_EQ(x_dist_attr.annotated().empty(), true);
|
||||
|
||||
std::stringstream x_sstream;
|
||||
x_sstream << x_dist_attr;
|
||||
EXPECT_EQ(x_sstream.str(), x_dist_attr.to_string());
|
||||
auto x_proto = phi::distributed::to_proto(x_dist_attr);
|
||||
TensorDistAttr new_x_dist_attr = get_dist_attr(x);
|
||||
new_x_dist_attr.from_proto(x_proto);
|
||||
EXPECT_EQ(x_dist_attr, new_x_dist_attr);
|
||||
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 0}));
|
||||
y_dist_attr.set_batch_dim(-1);
|
||||
y_dist_attr.set_chunk_id(0);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>({false, true}));
|
||||
x_dist_attr.mark_annotated("batch_dim");
|
||||
x_dist_attr.mark_annotated("dynamic_dims");
|
||||
EXPECT_EQ(y_dist_attr.process_mesh(), process_mesh);
|
||||
EXPECT_EQ(y_dist_attr.dims_mapping(), std::vector<int64_t>({-1, 0}));
|
||||
EXPECT_EQ(y_dist_attr.batch_dim(), -1);
|
||||
EXPECT_EQ(y_dist_attr.chunk_id(), 0);
|
||||
EXPECT_EQ(y_dist_attr.dynamic_dims(), std::vector<bool>({false, true}));
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("batch_dim"), true);
|
||||
EXPECT_EQ(x_dist_attr.is_annotated("dynamic_dims"), true);
|
||||
EXPECT_EQ(x_dist_attr.verify(get_tensor_shape(y)), true);
|
||||
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1}));
|
||||
out_dist_attr.set_batch_dim(1);
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
EXPECT_EQ(out_dist_attr.process_mesh(), process_mesh);
|
||||
EXPECT_EQ(out_dist_attr.dims_mapping(), std::vector<int64_t>({0, 1}));
|
||||
EXPECT_EQ(out_dist_attr.batch_dim(), 1);
|
||||
EXPECT_EQ(out_dist_attr.dynamic_dims(), std::vector<bool>({false, false}));
|
||||
EXPECT_EQ(out_dist_attr.verify(get_tensor_shape(out)), true);
|
||||
|
||||
OperatorDistAttr mul_dist_attr(*op);
|
||||
EXPECT_EQ(mul_dist_attr.impl_type(),
|
||||
paddle::distributed::auto_parallel::kDefault);
|
||||
EXPECT_EQ(mul_dist_attr.impl_idx(), 0);
|
||||
EXPECT_EQ(mul_dist_attr.chunk_id(), 0);
|
||||
EXPECT_EQ(mul_dist_attr.is_recompute(), false);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("process_mesh"), false);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_type"), false);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_idx"), false);
|
||||
mul_dist_attr.set_input_dist_attr(x->Name(), x_dist_attr);
|
||||
mul_dist_attr.set_input_dist_attr(y->Name(), y_dist_attr);
|
||||
mul_dist_attr.set_output_dist_attr(out->Name(), out_dist_attr);
|
||||
mul_dist_attr.set_process_mesh(process_mesh2);
|
||||
mul_dist_attr.set_impl_type("dist_mul");
|
||||
mul_dist_attr.set_impl_idx(0);
|
||||
mul_dist_attr.set_chunk_id(1);
|
||||
mul_dist_attr.set_is_recompute(true);
|
||||
mul_dist_attr.mark_annotated("process_mesh");
|
||||
mul_dist_attr.mark_annotated("impl_type");
|
||||
mul_dist_attr.mark_annotated("impl_idx");
|
||||
EXPECT_NE(mul_dist_attr.input_dist_attr(x->Name()), x_dist_attr);
|
||||
EXPECT_NE(mul_dist_attr.input_dist_attr(y->Name()), y_dist_attr);
|
||||
EXPECT_NE(mul_dist_attr.output_dist_attr(out->Name()), out_dist_attr);
|
||||
EXPECT_EQ(mul_dist_attr.process_mesh(), process_mesh2);
|
||||
EXPECT_EQ(mul_dist_attr.input_dist_attr(x->Name()).process_mesh(),
|
||||
process_mesh2);
|
||||
EXPECT_EQ(mul_dist_attr.input_dist_attr(y->Name()).process_mesh(),
|
||||
process_mesh2);
|
||||
EXPECT_EQ(mul_dist_attr.impl_type(), "dist_mul");
|
||||
EXPECT_EQ(mul_dist_attr.impl_idx(), 0);
|
||||
EXPECT_EQ(mul_dist_attr.chunk_id(), 1);
|
||||
EXPECT_EQ(mul_dist_attr.is_recompute(), true);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("process_mesh"), true);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_type"), true);
|
||||
EXPECT_EQ(mul_dist_attr.is_annotated("impl_idx"), true);
|
||||
EXPECT_EQ(mul_dist_attr.verify(op), true);
|
||||
mul_dist_attr.clear_annotated();
|
||||
EXPECT_EQ(mul_dist_attr.annotated().empty(), true);
|
||||
|
||||
std::stringstream mul_sstream;
|
||||
mul_sstream << mul_dist_attr;
|
||||
EXPECT_EQ(mul_sstream.str(), mul_dist_attr.to_string());
|
||||
auto mul_proto = mul_dist_attr.to_proto();
|
||||
OperatorDistAttr new_mul_dist_attr(*op);
|
||||
new_mul_dist_attr.from_proto(mul_proto);
|
||||
EXPECT_EQ(mul_dist_attr, new_mul_dist_attr);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,73 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "paddle/phi/core/distributed/auto_parallel/dist_mapper.h"
|
||||
#include <map>
|
||||
#include <sstream>
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(DistributedMapper, Ctor) {
|
||||
std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int64_t> device_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
std::string device_type = "GPU";
|
||||
int64_t size = shape[0] * shape[1];
|
||||
|
||||
DeviceMesh device_mesh("device_mesh", shape, device_ids, dim_names);
|
||||
for (int64_t i = 0; i < shape[0]; ++i) {
|
||||
for (int64_t j = 0; j < shape[1]; ++j) {
|
||||
int64_t global_id = i * shape[1] + j;
|
||||
int64_t local_id = j;
|
||||
int64_t machine_id = i;
|
||||
device_mesh.add_device(
|
||||
Device(global_id, local_id, machine_id, device_type));
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < size; ++i) {
|
||||
for (int64_t j = 0; j < size; ++j) {
|
||||
device_mesh.add_link(Link(i, j, "NVL"));
|
||||
}
|
||||
}
|
||||
|
||||
DistributedMapper dist_mapper;
|
||||
dist_mapper.add_device_mesh(device_mesh);
|
||||
std::map<int64_t, std::pair<std::string, std::vector<int64_t>>>
|
||||
process_id_to_device_ids;
|
||||
process_id_to_device_ids[0] = {"device_mesh", {5}};
|
||||
process_id_to_device_ids[1] = {"device_mesh", {4}};
|
||||
process_id_to_device_ids[2] = {"device_mesh", {3}};
|
||||
process_id_to_device_ids[3] = {"device_mesh", {2}};
|
||||
process_id_to_device_ids[4] = {"device_mesh", {1}};
|
||||
process_id_to_device_ids[5] = {"device_mesh", {0}};
|
||||
dist_mapper.set_process_id_to_device_ids(process_id_to_device_ids);
|
||||
|
||||
EXPECT_EQ(dist_mapper.device_meshes().at("device_mesh"), device_mesh);
|
||||
EXPECT_EQ(dist_mapper.device_mesh("device_mesh"), device_mesh);
|
||||
EXPECT_EQ(dist_mapper.process_id_to_device_ids(), process_id_to_device_ids);
|
||||
std::stringstream sstream;
|
||||
sstream << dist_mapper;
|
||||
EXPECT_EQ(sstream.str(), dist_mapper.to_string());
|
||||
auto proto = phi::distributed::to_proto(dist_mapper);
|
||||
DistributedMapper new_dist_mapper = DistributedMapper::from_proto(proto);
|
||||
EXPECT_EQ(dist_mapper, new_dist_mapper);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,72 @@
|
||||
/* Copyright (c) 2023 PaddlePaddle 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 "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
|
||||
#include "test/cpp/phi/core/allocator.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace tests {
|
||||
|
||||
TEST(dist_tensor, constructor) {
|
||||
auto fancy_allocator =
|
||||
std::unique_ptr<Allocator>(new phi::tests::FancyAllocator);
|
||||
auto* alloc = fancy_allocator.get();
|
||||
|
||||
DataType dtype{DataType::FLOAT16};
|
||||
DDim dims({3, 4});
|
||||
DenseTensorMeta meta(dtype, dims);
|
||||
|
||||
auto dist_attr = TensorDistAttr(common::vectorize(dims));
|
||||
|
||||
std::vector<int64_t> mesh_shape = {1};
|
||||
std::vector<int64_t> process_ids = {0};
|
||||
std::vector<std::string> dim_names = {"x"};
|
||||
ProcessMesh mesh(mesh_shape, process_ids, dim_names);
|
||||
dist_attr.set_process_mesh(mesh);
|
||||
|
||||
// copy construct
|
||||
std::shared_ptr<DenseTensor> x1 = std::make_shared<DenseTensor>(alloc, meta);
|
||||
DistTensor dist_x1(x1, dist_attr);
|
||||
EXPECT_TRUE(dist_x1.defined());
|
||||
EXPECT_TRUE(dist_x1.initialized());
|
||||
EXPECT_TRUE(dist_x1.valid());
|
||||
EXPECT_EQ(dist_x1.numel(), 12L);
|
||||
EXPECT_EQ(dist_x1.local_dims()[0], 3L);
|
||||
EXPECT_EQ(dist_x1.local_dims()[1], 4L);
|
||||
|
||||
// empty construct
|
||||
DistTensor dist_x2(dims, dist_attr);
|
||||
EXPECT_TRUE(!dist_x2.defined());
|
||||
EXPECT_TRUE(!dist_x2.initialized());
|
||||
// allocate error test
|
||||
bool caught_exception = false;
|
||||
try {
|
||||
dist_x2.AllocateFrom(alloc, phi::DataType::FLOAT32, 12L, false);
|
||||
} catch (common::enforce::EnforceNotMet& error) {
|
||||
caught_exception = true;
|
||||
EXPECT_NE(std::string(error.what()).find("Unavailable"), 0UL);
|
||||
}
|
||||
EXPECT_TRUE(caught_exception);
|
||||
}
|
||||
|
||||
} // namespace tests
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,116 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(ExpandAsInferSpmd, Ctor) {
|
||||
// Sharding along axes besides softmax axis.
|
||||
std::vector<int64_t> x_shape = {1, 48};
|
||||
std::vector<int64_t> y_shape = {2, 32, 48};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
TensorDistAttr y_dist_attr = TensorDistAttr();
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor y(phi::make_ddim(y_shape), y_dist_attr);
|
||||
|
||||
// test info forward
|
||||
auto spmdinfo = ExpandAsInferSpmd(x, y, y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test ExpandAsInferSpmd" << std::endl << std::endl << std::endl;
|
||||
|
||||
// test info forward without y
|
||||
// x [1, 48], target [2, 32, 48]: [-1, -1] -> [-1, -1, -1]
|
||||
spmdinfo = ExpandAsInferSpmd(x, phi::distributed::DistMetaTensor(), y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
const phi::distributed::ArgDistAttr& attr = spmdinfo.first[1];
|
||||
if (paddle::holds_alternative<phi::distributed::TensorDistAttr>(attr)) {
|
||||
EXPECT_EQ(paddle::get<phi::distributed::TensorDistAttr>(attr),
|
||||
phi::distributed::TensorDistAttr());
|
||||
} else {
|
||||
FAIL() << "forward_info.first[1] is not TensorDistAttr";
|
||||
}
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({-1, -1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test ExpandAsInferSpmd" << std::endl << std::endl << std::endl;
|
||||
|
||||
// test info reverse
|
||||
spmdinfo = ExpandAsInferSpmdReverse(x, y, y, y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test ExpandAsInferSpmdReverse" << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// test info grad
|
||||
spmdinfo = ExpandAsGradInferSpmd(x, y, y_shape);
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
check_partial_dims(spmdinfo.second[0], {0, 1});
|
||||
VLOG(4) << "Test ExpandAsGradInferSpmd" << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,105 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 "glog/logging.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
ProcessMesh CreateProcessMesh() {
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
return ProcessMesh(mesh_shape, process_ids, dim_names);
|
||||
}
|
||||
|
||||
phi::distributed::DistMetaTensor CreateDistMetaTensor(
|
||||
const std::vector<int64_t>& shape,
|
||||
const std::vector<int64_t>& dims_mapping,
|
||||
const ProcessMesh& process_mesh) {
|
||||
TensorDistAttr dist_attr;
|
||||
dist_attr.set_process_mesh(process_mesh);
|
||||
dist_attr.set_dims_mapping(dims_mapping);
|
||||
return phi::distributed::DistMetaTensor(phi::make_ddim(shape), dist_attr);
|
||||
}
|
||||
|
||||
TEST(ExpandInferSpmd, Ctor) {
|
||||
ProcessMesh process_mesh = CreateProcessMesh();
|
||||
|
||||
// Test case forward 1: Expand with shape {8, 2, 6, 1024, -1}
|
||||
auto x = CreateDistMetaTensor(
|
||||
{8, 2, 1, 1024, 128}, {0, -1, -1, 1, -1}, process_mesh);
|
||||
phi::IntArray shape = {8, 2, 6, 1024, -1};
|
||||
auto spmdinfo = ExpandInferSpmd(x, shape);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({0, -1, -1, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, -1, -1, 1, -1}));
|
||||
|
||||
// Test case forward 2: Expand with shape {2, -1}
|
||||
auto x1 = CreateDistMetaTensor({8}, {1}, process_mesh);
|
||||
phi::IntArray shape1 = {2, -1};
|
||||
auto spmdinfo1 = ExpandInferSpmd(x1, shape1);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo1.first[0]), std::vector<int64_t>({1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo1.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
|
||||
// Test case forward 3: Expand with shape {0, -1}
|
||||
auto x2 = CreateDistMetaTensor({8}, {1}, process_mesh);
|
||||
phi::IntArray shape2 = {0, -1};
|
||||
auto spmdinfo2 = ExpandInferSpmd(x2, shape2);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo2.first[0]), std::vector<int64_t>({1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo2.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
|
||||
// Test case backward 1: ExpandGrad with shape {0, -1}
|
||||
auto x3 = CreateDistMetaTensor({8}, {1}, process_mesh);
|
||||
auto out3 = CreateDistMetaTensor({2, 8}, {-1, 1}, process_mesh);
|
||||
phi::IntArray shape3 = {0, -1};
|
||||
auto spmdinfo3 = ExpandGradInferSpmd(x3, out3, shape3);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo3.first[0]), std::vector<int64_t>({1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo3.first[1]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo3.second[0]), std::vector<int64_t>({1}));
|
||||
|
||||
// Test case backward 2: ExpandGrad with shape {2, 2, -1}
|
||||
auto x4 = CreateDistMetaTensor({1, 8}, {-1, 1}, process_mesh);
|
||||
auto out4 = CreateDistMetaTensor({2, 2, 8}, {-1, -1, 1}, process_mesh);
|
||||
phi::IntArray shape4 = {2, 2, -1};
|
||||
auto spmdinfo4 = ExpandGradInferSpmd(x4, out4, shape4);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo4.first[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo4.first[1]),
|
||||
std::vector<int64_t>({-1, -1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo4.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
|
||||
// Test case backward 3: ExpandGrad with shape {2, 2, -1}
|
||||
auto x5 = CreateDistMetaTensor({1, 8}, {-1, 1}, process_mesh);
|
||||
auto out5 = CreateDistMetaTensor({2, 2, 8}, {-1, 0, 1}, process_mesh);
|
||||
phi::IntArray shape5 = {2, 2, -1};
|
||||
auto spmdinfo5 = ExpandGradInferSpmd(x5, out5, shape5);
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo5.first[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo5.first[1]),
|
||||
std::vector<int64_t>({-1, 0, 1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo5.second[0]),
|
||||
std::vector<int64_t>({-1, 1}));
|
||||
EXPECT_EQ(get_partial_dims(spmdinfo5.second[0]), std::set<int64_t>({0}));
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,103 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(FusedLinearParamGradAddSPMDRule, Ctor) {
|
||||
// build input data class
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 4};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
// b s h
|
||||
std::vector<int64_t> x_shape = {2, 512, 1024};
|
||||
std::vector<int64_t> out_shape = {2, 512, 2048};
|
||||
std::vector<int64_t> weight_shape = {1024, 2048};
|
||||
std::vector<int64_t> bias_shape = {2048};
|
||||
|
||||
// test mp col split
|
||||
{
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor out(phi::make_ddim(out_shape),
|
||||
out_dist_attr);
|
||||
phi::distributed::DistMetaTensor dweight;
|
||||
phi::distributed::DistMetaTensor dbias;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto spmd_info =
|
||||
FusedLinearParamGradAddInferSpmd(x, out, dweight, dbias, 0, true);
|
||||
check_dim_mapping(spmd_info.second[0], {-1, 1});
|
||||
check_partial_dims(spmd_info.second[0], {0});
|
||||
check_dim_mapping(spmd_info.second[1], {1});
|
||||
check_partial_dims(spmd_info.second[1], {0});
|
||||
dweight = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(weight_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[0]));
|
||||
dbias = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(bias_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[1]));
|
||||
}
|
||||
}
|
||||
|
||||
// test mp row split
|
||||
{
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, 1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1}));
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor out(phi::make_ddim(out_shape),
|
||||
out_dist_attr);
|
||||
phi::distributed::DistMetaTensor dweight;
|
||||
phi::distributed::DistMetaTensor dbias;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto spmd_info =
|
||||
FusedLinearParamGradAddInferSpmd(x, out, dweight, dbias, 0, true);
|
||||
check_dim_mapping(spmd_info.second[0], {1, -1});
|
||||
check_partial_dims(spmd_info.second[0], {0});
|
||||
check_dim_mapping(spmd_info.second[1], {-1});
|
||||
check_partial_dims(spmd_info.second[1], {0});
|
||||
dweight = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(weight_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[0]));
|
||||
dbias = phi::distributed::DistMetaTensor(
|
||||
phi::make_ddim(bias_shape),
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmd_info.second[1]));
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,110 @@
|
||||
/* Copyright (c) 2024 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
TEST(FusedRmsNormSPMDRule, test_fused_rms_norm) {
|
||||
// build input data class
|
||||
std::vector<int64_t> x_shape = {64, 32, 1024};
|
||||
std::vector<int64_t> scale_shape = {1024};
|
||||
std::vector<int64_t> variance_shape = {64, 32};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
|
||||
TensorDistAttr scale_dist_attr = TensorDistAttr();
|
||||
scale_dist_attr.set_process_mesh(process_mesh);
|
||||
scale_dist_attr.set_dims_mapping(std::vector<int64_t>({-1}));
|
||||
scale_dist_attr.set_dynamic_dims(std::vector<bool>({false}));
|
||||
|
||||
x_dist_attr.set_dims_mapping({1, -1, -1});
|
||||
scale_dist_attr.set_dims_mapping(std::vector<int64_t>{-1});
|
||||
|
||||
phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor scale(common::make_ddim(scale_shape),
|
||||
scale_dist_attr);
|
||||
auto inferred_dist_attrs = phi::distributed::RmsNormInferSpmd(x, scale, 0.5);
|
||||
|
||||
size_t input_size = 2;
|
||||
size_t output_size = 2;
|
||||
EXPECT_EQ(inferred_dist_attrs.first.size(), input_size);
|
||||
EXPECT_EQ(inferred_dist_attrs.second.size(), output_size);
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {1, -1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {1, -1});
|
||||
|
||||
VLOG(4) << "test1 done.";
|
||||
|
||||
x_dist_attr.set_dims_mapping({1, 0, -1});
|
||||
scale_dist_attr.set_dims_mapping(std::vector<int64_t>{0});
|
||||
x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
|
||||
scale = phi::distributed::DistMetaTensor(common::make_ddim(scale_shape),
|
||||
scale_dist_attr);
|
||||
|
||||
inferred_dist_attrs = phi::distributed::RmsNormInferSpmd(x, scale, 0.5);
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {1, 0, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {1, 0, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {1, 0});
|
||||
VLOG(4) << "test2 done.";
|
||||
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1}));
|
||||
out_dist_attr.set_dynamic_dims(std::vector<bool>({false, false, false}));
|
||||
phi::distributed::DistMetaTensor out(common::make_ddim(x_shape),
|
||||
out_dist_attr);
|
||||
|
||||
TensorDistAttr invvar_dist_attr = TensorDistAttr();
|
||||
invvar_dist_attr.set_process_mesh(process_mesh);
|
||||
invvar_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1}));
|
||||
invvar_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
phi::distributed::DistMetaTensor invvar(common::make_ddim(variance_shape),
|
||||
invvar_dist_attr);
|
||||
|
||||
inferred_dist_attrs =
|
||||
phi::distributed::RmsNormInferSpmdReverse(x, scale, out, invvar, 0.5);
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {0, 1});
|
||||
VLOG(4) << "test3 done.";
|
||||
|
||||
x_dist_attr.set_dims_mapping({0, 1, -1});
|
||||
x = phi::distributed::DistMetaTensor(common::make_ddim(x_shape), x_dist_attr);
|
||||
inferred_dist_attrs =
|
||||
phi::distributed::RmsNormGradInferSpmd(x, scale, invvar, out, 0.5);
|
||||
|
||||
check_dim_mapping(inferred_dist_attrs.first[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[1], {-1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[2], {0, 1});
|
||||
check_dim_mapping(inferred_dist_attrs.first[3], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[0], {0, 1, -1});
|
||||
check_dim_mapping(inferred_dist_attrs.second[1], {-1});
|
||||
check_partial_dims(inferred_dist_attrs.second[1], {0, 1});
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,286 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 <set>
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct IndexSelectTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
std::vector<int64_t> index_shape;
|
||||
std::vector<std::vector<int64_t>> index_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_index_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
};
|
||||
|
||||
struct IndexSelectGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
std::vector<int64_t> index_shape;
|
||||
std::vector<std::vector<int64_t>> index_dims_mapping;
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_index_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
std::set<int64_t> partial_dims;
|
||||
};
|
||||
|
||||
TEST(IndexSelectInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<IndexSelectTestCase> test_cases = {
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0,1],[2],[]], [[]] -> [[0,1],[],[]], [[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0,1],[2],[]], [[2]] -> [[0,1],[],[]], [[2]], [[0,1],[2],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{2}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{2}},
|
||||
{{0, 1}, {2}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0,1],[2],[]], [[0]] -> [[0,1],[],[]], [[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{0}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[2],[],[]], [[0,1]] -> [[2],[],[]], [[0,1]], [[2],[0,1],[]]
|
||||
{{8, 16, 32},
|
||||
{{2}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
1,
|
||||
{{2}, {}, {}},
|
||||
{{0, 1}},
|
||||
{{2}, {0, 1}, {}}},
|
||||
|
||||
// [8, 16, 32], [8], axis = 1
|
||||
// [[0],[],[]], [[0,1]] -> [[0],[],[]], [[1]], [[0],[1],[]]
|
||||
{{8, 16, 32},
|
||||
{{0}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
1,
|
||||
{{0}, {}, {}},
|
||||
{{1}},
|
||||
{{0}, {1}, {}}},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr index_dist_attr = TensorDistAttr();
|
||||
index_dist_attr.set_process_mesh(process_mesh);
|
||||
index_dist_attr.set_dims_mapping(tc.index_dims_mapping);
|
||||
index_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.index_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor index = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.index_shape), index_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::IndexSelectInferSpmd(x, index, tc.axis);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.first[1],
|
||||
tc.expected_index_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(IndexSelectGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<IndexSelectGradTestCase> test_cases = {
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0,1],[2],[]], [[]], [[0,1], [], []] -> [[0,1],[],[]], [[]],
|
||||
// [[0,1],[],[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{}},
|
||||
{8, 8, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0,1],[2],[]], [[2]], [[0,1],[2],[]] -> [[0,1],[],[]], [[2]],
|
||||
// [[0,1],[2],[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{2}},
|
||||
{8, 8, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{2}},
|
||||
{{0, 1}, {2}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{2}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0,1],[2],[]], [[0]], [[0,1],[],[]] -> [[0,1],[],[]], [[]],
|
||||
// [[0,1],[],[]], [[0,1],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0, 1}, {2}, {}},
|
||||
{8},
|
||||
{{0}},
|
||||
{8, 8, 32},
|
||||
{{0, 1}, {}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[2],[],[]], [[0,1]], [[2],[0,1],[]] -> [[2],[],[]], [[0,1]],
|
||||
// [[2],[0,1],[]], [[2],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{2}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
{8, 8, 32},
|
||||
{{2}, {0, 1}, {}},
|
||||
1,
|
||||
{{2}, {}, {}},
|
||||
{{0, 1}},
|
||||
{{2}, {0, 1}, {}},
|
||||
{{2}, {}, {}},
|
||||
{0, 1}},
|
||||
|
||||
// [8, 16, 32], [8], [8, 8, 32], axis = 1
|
||||
// [[0],[],[]], [[0,1]], [[0],[1],[]] -> [[0],[],[]], [[1]], [[0],[1],[]],
|
||||
// [[0],[],[]]
|
||||
{{8, 16, 32},
|
||||
{{0}, {}, {}},
|
||||
{8},
|
||||
{{0, 1}},
|
||||
{8, 8, 32},
|
||||
{{0}, {1}, {}},
|
||||
1,
|
||||
{{0}, {}, {}},
|
||||
{{1}},
|
||||
{{0}, {1}, {}},
|
||||
{{0}, {}, {}},
|
||||
{1}},
|
||||
};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr index_dist_attr = TensorDistAttr();
|
||||
index_dist_attr.set_process_mesh(process_mesh);
|
||||
index_dist_attr.set_dims_mapping(tc.index_dims_mapping);
|
||||
index_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.index_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor index = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.index_shape), index_dist_attr);
|
||||
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::IndexSelectGradInferSpmd(x, index, out_grad, tc.axis);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_index_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[2],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
if (!tc.partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(backward_spmd_info.second[0]), true);
|
||||
check_partial_dims(backward_spmd_info.second[0], tc.partial_dims);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,505 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 <set>
|
||||
#include "paddle/phi/infermeta/spmd_rules/bmm.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct MatmulTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<int64_t> y_shape;
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping;
|
||||
|
||||
// attribute
|
||||
bool trans_x;
|
||||
bool trans_y;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_y_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
|
||||
std::set<int64_t> partial_dims;
|
||||
};
|
||||
|
||||
struct MatmulGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<int64_t> y_shape;
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping;
|
||||
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// attribute
|
||||
bool trans_x;
|
||||
bool trans_y;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_y_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_y_grad_dims_mapping;
|
||||
|
||||
std::set<int64_t> x_grad_partial_dims;
|
||||
std::set<int64_t> y_grad_partial_dims;
|
||||
};
|
||||
|
||||
TEST(MatmulInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<MatmulTestCase> test_cases = {
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1], []] ,[[],[2]] -> [[0,1], []] ,[[],[2]],[[0,1],[2]]
|
||||
{{64, 32},
|
||||
{{0, 1}, {}},
|
||||
{32, 48},
|
||||
{{}, {2}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {}},
|
||||
{{}, {2}},
|
||||
{{0, 1}, {2}},
|
||||
{}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1], [2]] ,[[],[]] -> [[0,1], [2]] ,[[2],[]],[[0,1],[]], partial: 2
|
||||
{{64, 32},
|
||||
{{0, 1}, {2}},
|
||||
{32, 48},
|
||||
{{}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {2}},
|
||||
{{2}, {}},
|
||||
{{0, 1}, {}},
|
||||
{2}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[], []] ,[[0,1],[2]] -> [[],[0,1]] ,[[0,1],[2],[[],[2]], partial:
|
||||
// {0,1}
|
||||
{{64, 32},
|
||||
{{}, {}},
|
||||
{32, 48},
|
||||
{{0, 1}, {2}},
|
||||
false,
|
||||
false,
|
||||
{{}, {0, 1}},
|
||||
{{0, 1}, {2}},
|
||||
{{}, {2}},
|
||||
{0, 1}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0], [1]] ,[[2],[0]] -> [[0], [1,2]] ,[[1,2],[]],[[0],[]], partial:
|
||||
// {1,2}
|
||||
{{64, 32},
|
||||
{{0}, {1}},
|
||||
{32, 48},
|
||||
{{2}, {0}},
|
||||
false,
|
||||
false,
|
||||
{{0}, {1, 2}},
|
||||
{{1, 2}, {}},
|
||||
{{0}, {}},
|
||||
{1, 2}},
|
||||
|
||||
// [64, 32], [32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1], [2]] ,[[0],[]] -> [[0,1], [2]] ,[[2],[]],[[0,1],[]], partial: 2
|
||||
{{64, 32},
|
||||
{{0, 1}, {2}},
|
||||
{32, 48},
|
||||
{{0}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {2}},
|
||||
{{2}, {}},
|
||||
{{0, 1}, {}},
|
||||
{2}},
|
||||
|
||||
// [512, 48, 64, 32], [1, 32, 48], trans_x=false, trans_y=false
|
||||
// [[0,1],[2],[],[]] ,[[],[],[]] -> [[0,1],[2],[],[]]
|
||||
// ,[[],[],[]],[[0,1],[2],[],[]],
|
||||
// partial: {}
|
||||
{{512, 48, 64, 32},
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
{{}, {}, {}},
|
||||
{{0, 1}, {2}, {}, {}},
|
||||
{}},
|
||||
|
||||
// [512, 48, 32, 64], [1, 32, 48], trans_x=true, trans_y=false
|
||||
// [[0],[],[1,2],[]] ,[[],[],[2]] -> [[0],[],[1],[]]
|
||||
// ,[[],[1],[2]],[[0],[],[],[2]],
|
||||
// partial: {1}
|
||||
{{512, 48, 32, 64},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
true,
|
||||
false,
|
||||
{{0}, {}, {1}, {}},
|
||||
{{}, {1}, {2}},
|
||||
{{0}, {}, {}, {2}},
|
||||
{1}},
|
||||
|
||||
// [512, 48, 64, 32], [1, 48, 32], trans_x=false, trans_y=true
|
||||
// [[0],[],[1,2],[]] ,[[],[0],[]] -> [[0],[],[1,2],[]]
|
||||
// ,[[],[],[]],[[0],[],[1,2],[]],
|
||||
// partial: {}
|
||||
{{512, 48, 64, 32},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{1, 48, 32},
|
||||
{{}, {0}, {}},
|
||||
false,
|
||||
true,
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{}},
|
||||
|
||||
// [512, 48, 32, 64], [1, 48, 32], trans_x=true, trans_y=true
|
||||
// [[],[],[0,1],[2]] ,[[],[0,1],[2]] -> [[],[],[],[2]]
|
||||
// ,[[],[0,1],[]],[[],[],[2],[0,1]],
|
||||
// partial: {}
|
||||
{{512, 48, 32, 64},
|
||||
{{}, {}, {0, 1}, {2}},
|
||||
{1, 48, 32},
|
||||
{{}, {0, 1}, {2}},
|
||||
true,
|
||||
true,
|
||||
{{}, {}, {}, {2}},
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {}, {2}, {0, 1}},
|
||||
{}},
|
||||
};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr y_dist_attr = TensorDistAttr();
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(tc.y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(tc.y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.y_shape), y_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::MatmulInferSpmd(x, y, tc.trans_x, tc.trans_y);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.first[1],
|
||||
tc.expected_y_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
if (!tc.partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(forward_spmd_info.second[0]), true);
|
||||
check_partial_dims(forward_spmd_info.second[0], tc.partial_dims);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(MatmulGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<MatmulGradTestCase> test_cases = {
|
||||
// [64, 32], [32, 48], [64,48], trans_x=false, trans_y=false
|
||||
// [[0,1], []] ,[[],[2]], [[0,1],[2]] -> [[0,1], []]
|
||||
// ,[[],[2]],[[0,1],[2]], [[0,1],[]], [[],[2]], x_partial: {2}, y_partial:
|
||||
// {0,1}
|
||||
{{64, 32},
|
||||
{{0, 1}, {}},
|
||||
{32, 48},
|
||||
{{}, {2}},
|
||||
{64, 48},
|
||||
{{0, 1}, {2}},
|
||||
false,
|
||||
false,
|
||||
{{0, 1}, {}},
|
||||
{{}, {2}},
|
||||
{{0, 1}, {2}},
|
||||
{{0, 1}, {}},
|
||||
{{}, {2}},
|
||||
{2},
|
||||
{0, 1}},
|
||||
// [1024,512,64,32], [1,32,48], [1024,512,64,48], trans_x=false,
|
||||
// trans_y=false
|
||||
// [[0],[],[1,2],[]] ,[[],[],[2]], [[0],[],[1,2],[]] -> [[0],[],[1,2],[]]
|
||||
// ,[[],[],[]], [[0],[],[1,2],[]], [[0],[],[1,2],[]], [[],[],[]],
|
||||
// x_grad_partial: {}, y_grad_partial: {0,1,2}
|
||||
{{1024, 512, 64, 32},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
{1024, 512, 64, 48},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
false,
|
||||
false,
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{0}, {}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{},
|
||||
{0, 1, 2}},
|
||||
// [1024,512,64,32], [1,32,48], [1024,512,64,48], trans_x=false,
|
||||
// trans_y=false
|
||||
// [[],[0],[1,2],[]] ,[[],[],[2]], [[],[0],[1,2],[]] -> [[],[0],[1,2],[]]
|
||||
// ,[[],[],[]], [[],[0],[1,2],[]], [[],[0],[1,2],[]], [[],[],[]],
|
||||
// x_grad_partial: {}, y_grad_partial: {0,1,2}
|
||||
{{1024, 512, 64, 32},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
false,
|
||||
false,
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{{}, {}, {}},
|
||||
{},
|
||||
{0, 1, 2}},
|
||||
// [1024,512,32,64], [1,32,48], [1024,512,64,48], trans_x=true,
|
||||
// trans_y=false
|
||||
// [[],[0],[1,2],[]] ,[[],[],[2]], [[],[0],[],[2]] -> [[],[0],[1],[]]
|
||||
// ,[[],[1],[2]], [[],[0],[],[2]], [[],[0],[1],[]], [[],[1],[2]],
|
||||
// x_grad_partial: {2}, y_grad_partial: {0}
|
||||
{{1024, 512, 32, 64},
|
||||
{{}, {0}, {1, 2}, {}},
|
||||
{1, 32, 48},
|
||||
{{}, {}, {2}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {0}, {}, {2}},
|
||||
true,
|
||||
false,
|
||||
{{}, {0}, {1}, {}},
|
||||
{{}, {1}, {2}},
|
||||
{{}, {0}, {}, {2}},
|
||||
{{}, {0}, {1}, {}},
|
||||
{{}, {1}, {2}},
|
||||
{2},
|
||||
{0}},
|
||||
// [1024,512,32,64], [1,48,32], [1024,512,64,48], trans_x=true,
|
||||
// trans_y=true
|
||||
// [[],[],[1,2],[]] ,[[],[],[0]], [[],[],[],[]] -> [[],[],[0,1,2],[]]
|
||||
// ,[[],[],[0,1,2]], [[],[],[],[]], [[],[],[0,1,2],[]], [[],[],[0,1,2]],
|
||||
// x_grad_partial: {}, y_grad_partial: {}
|
||||
{{1024, 512, 32, 64},
|
||||
{{}, {}, {1, 2}, {}},
|
||||
{1, 48, 32},
|
||||
{{}, {}, {0}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {}, {}, {}},
|
||||
true,
|
||||
true,
|
||||
{{}, {}, {1, 2, 0}, {}},
|
||||
{{}, {}, {1, 2, 0}},
|
||||
{{}, {}, {}, {}},
|
||||
{{}, {}, {1, 2, 0}, {}},
|
||||
{{}, {}, {1, 2, 0}},
|
||||
{},
|
||||
{}},
|
||||
// [1024,512,64,32], [1,48,32], [1024,512,64,48], trans_x=false,
|
||||
// trans_y=true
|
||||
// [[],[],[0],[1,2]] ,[[],[],[0]], [[],[],[0],[]] -> [[],[],[0],[1,2]]
|
||||
// ,[[],[],[1,2]], [[],[],[0],[]], [[],[],[0],[1,2]],
|
||||
// [[],[],[1,2]],
|
||||
// x_grad_partial: {}, y_grad_partial: {0}
|
||||
{{1024, 512, 64, 32},
|
||||
{{}, {}, {0}, {1, 2}},
|
||||
{1, 48, 32},
|
||||
{{}, {}, {0}},
|
||||
{1024, 512, 64, 48},
|
||||
{{}, {}, {0}, {}},
|
||||
false,
|
||||
true,
|
||||
{{}, {}, {0}, {1, 2}},
|
||||
{{}, {}, {1, 2}},
|
||||
{{}, {}, {0}, {}},
|
||||
{{}, {}, {0}, {1, 2}},
|
||||
{{}, {}, {1, 2}},
|
||||
{},
|
||||
{0}}};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
TensorDistAttr y_dist_attr = TensorDistAttr();
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(tc.y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(tc.y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.y_shape), y_dist_attr);
|
||||
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::MatmulGradInferSpmd(
|
||||
x, y, out_grad, tc.trans_x, tc.trans_y);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(2));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_y_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[2],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
if (!tc.x_grad_partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(backward_spmd_info.second[0]), true);
|
||||
check_partial_dims(backward_spmd_info.second[0], tc.x_grad_partial_dims);
|
||||
}
|
||||
check_multi_dims_mapping(backward_spmd_info.second[1],
|
||||
tc.expected_y_grad_dims_mapping);
|
||||
if (!tc.y_grad_partial_dims.empty()) {
|
||||
EXPECT_EQ(is_partial(backward_spmd_info.second[1]), true);
|
||||
check_partial_dims(backward_spmd_info.second[1], tc.y_grad_partial_dims);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(BmmInferSpmd, CoShard) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> x_shape = {4, 16, 8};
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping = {{0, 1}, {2}, {}};
|
||||
TensorDistAttr x_dist_attr;
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
|
||||
|
||||
std::vector<int64_t> y_shape = {4, 8, 32};
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping = {{0, 1}, {}, {}};
|
||||
TensorDistAttr y_dist_attr;
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
|
||||
|
||||
auto bmm_spmd_info = phi::distributed::BmmInferSpmd(x, y);
|
||||
|
||||
ASSERT_EQ(bmm_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
ASSERT_EQ(bmm_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
|
||||
check_multi_dims_mapping(bmm_spmd_info.first[0], x_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_spmd_info.first[0]));
|
||||
check_multi_dims_mapping(bmm_spmd_info.first[1], y_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_spmd_info.first[1]));
|
||||
|
||||
const std::vector<std::vector<int64_t>> expected_out_dims_mapping = {
|
||||
{0, 1}, {2}, {}};
|
||||
check_multi_dims_mapping(bmm_spmd_info.second[0], expected_out_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_spmd_info.second[0]));
|
||||
}
|
||||
|
||||
TEST(BmmGradInferSpmd, CoShard) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> x_shape = {4, 16, 8};
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping = {{0, 1}, {2}, {}};
|
||||
TensorDistAttr x_dist_attr;
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x(common::make_ddim(x_shape), x_dist_attr);
|
||||
|
||||
std::vector<int64_t> y_shape = {4, 8, 32};
|
||||
std::vector<std::vector<int64_t>> y_dims_mapping = {{0, 1}, {}, {}};
|
||||
TensorDistAttr y_dist_attr;
|
||||
y_dist_attr.set_process_mesh(process_mesh);
|
||||
y_dist_attr.set_dims_mapping(y_dims_mapping);
|
||||
y_dist_attr.set_dynamic_dims(std::vector<bool>(y_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor y(common::make_ddim(y_shape), y_dist_attr);
|
||||
|
||||
std::vector<int64_t> out_grad_shape = {4, 16, 32};
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping = {{0, 1}, {2}, {}};
|
||||
TensorDistAttr out_grad_dist_attr;
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(out_grad_dims_mapping);
|
||||
out_grad_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad(common::make_ddim(out_grad_shape),
|
||||
out_grad_dist_attr);
|
||||
|
||||
auto bmm_grad_spmd_info = phi::distributed::BmmGradInferSpmd(x, y, out_grad);
|
||||
|
||||
ASSERT_EQ(bmm_grad_spmd_info.first.size(), static_cast<size_t>(3));
|
||||
ASSERT_EQ(bmm_grad_spmd_info.second.size(), static_cast<size_t>(2));
|
||||
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.first[0], x_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.first[0]));
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.first[1], y_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.first[1]));
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.first[2], out_grad_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.first[2]));
|
||||
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.second[0], x_dims_mapping);
|
||||
EXPECT_FALSE(is_partial(bmm_grad_spmd_info.second[0]));
|
||||
check_multi_dims_mapping(bmm_grad_spmd_info.second[1], y_dims_mapping);
|
||||
EXPECT_TRUE(is_partial(bmm_grad_spmd_info.second[1]));
|
||||
check_partial_dims(bmm_grad_spmd_info.second[1], {2});
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,144 @@
|
||||
/* Copyright (c) 2024 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using phi::distributed::ArgDistAttr;
|
||||
using phi::distributed::DistMetaTensor;
|
||||
|
||||
void test_moe_combine_spmd(
|
||||
const std::vector<std::vector<int64_t>>& input_shapes,
|
||||
const std::vector<std::vector<int64_t>>& input_dims_mappings,
|
||||
const std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>& expected_dims_mappings,
|
||||
bool test_bwd_spmd = false) {
|
||||
size_t num_inputs = 0;
|
||||
if (test_bwd_spmd) {
|
||||
num_inputs = 4;
|
||||
} else {
|
||||
num_inputs = 3;
|
||||
}
|
||||
|
||||
EXPECT_EQ(input_shapes.size(), num_inputs)
|
||||
<< "The number of input_shapes must be" << num_inputs << ", but got "
|
||||
<< input_shapes.size();
|
||||
EXPECT_EQ(input_dims_mappings.size(), num_inputs)
|
||||
<< "The number of input_dims_mapping must be" << num_inputs
|
||||
<< ", but got " << input_dims_mappings.size();
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"dp", "mp", "pp"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<DistMetaTensor> dist_meta_tensors;
|
||||
for (size_t i = 0; i < num_inputs; ++i) {
|
||||
TensorDistAttr dist_attr = TensorDistAttr();
|
||||
dist_attr.set_process_mesh(process_mesh);
|
||||
|
||||
const std::vector<int64_t>& shape = input_shapes[i];
|
||||
const std::vector<int64_t>& dim_mapping = input_dims_mappings[i];
|
||||
EXPECT_EQ(shape.size(), dim_mapping.size())
|
||||
<< "The size of shape and dim_mapping for input " << i
|
||||
<< " must be the same, but got " << shape.size()
|
||||
<< " != " << dim_mapping.size();
|
||||
|
||||
dist_attr.set_dims_mapping(dim_mapping);
|
||||
dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
|
||||
|
||||
dist_meta_tensors.push_back(
|
||||
DistMetaTensor(common::make_ddim(shape), dist_attr));
|
||||
}
|
||||
|
||||
phi::distributed::SpmdInfo spmd_info;
|
||||
if (test_bwd_spmd) {
|
||||
spmd_info = phi::distributed::MoECombineGradInferSpmd(dist_meta_tensors[0],
|
||||
dist_meta_tensors[1],
|
||||
dist_meta_tensors[2],
|
||||
dist_meta_tensors[3]);
|
||||
} else {
|
||||
spmd_info = phi::distributed::MoECombineInferSpmd(
|
||||
dist_meta_tensors[0], dist_meta_tensors[1], dist_meta_tensors[2]);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; ++i) {
|
||||
std::vector<ArgDistAttr> dist_attrs;
|
||||
std::vector<std::vector<int64_t>> dims_mappings;
|
||||
if (i == 0) {
|
||||
dist_attrs = spmd_info.first;
|
||||
dims_mappings = expected_dims_mappings.first;
|
||||
} else {
|
||||
dist_attrs = spmd_info.second;
|
||||
dims_mappings = expected_dims_mappings.second;
|
||||
}
|
||||
EXPECT_EQ(dist_attrs.size(), dims_mappings.size())
|
||||
<< "The size of dist_attr and expected_dims must be the same, but got "
|
||||
<< dist_attrs.size() << " != " << dims_mappings.size();
|
||||
|
||||
for (size_t j = 0; j < dist_attrs.size(); ++j) {
|
||||
const ArgDistAttr& inferred_attr = dist_attrs[j];
|
||||
const std::vector<int64_t>& expected_dims_mapping = dims_mappings[j];
|
||||
check_dim_mapping(inferred_attr, expected_dims_mapping);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(MoECombineSPMDRule, test_moe_combine_spmd) {
|
||||
// forward: x, combine_weights, scatter_index -> y
|
||||
// backward: x, combine_weights, scatter_index, grad_y -> grad_x,
|
||||
// grad_combine_weights
|
||||
|
||||
int s = 1024, h = 512, k = 2;
|
||||
const std::vector<std::vector<int64_t>>& forward_input_shapes = {
|
||||
{s * k, h}, {s, k}, {s, k}};
|
||||
const std::vector<std::vector<int64_t>>& backward_input_shapes = {
|
||||
{s * k, h}, {s, k}, {s, k}, {s, h}};
|
||||
|
||||
// replicated case, forward
|
||||
std::vector<std::vector<int64_t>> input_dims_mappings = {
|
||||
{-1, -1}, {-1, -1}, {-1, -1}};
|
||||
std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>
|
||||
expected_dims_mappings = {{{-1, -1}, {-1, -1}, {-1, -1}}, {{-1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
forward_input_shapes, input_dims_mappings, expected_dims_mappings);
|
||||
|
||||
// replicated case, backward
|
||||
input_dims_mappings = {{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}};
|
||||
expected_dims_mappings = {{{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1}},
|
||||
{{-1, -1}, {-1, -1}, {-1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
backward_input_shapes, input_dims_mappings, expected_dims_mappings, true);
|
||||
|
||||
// mp case, forward
|
||||
input_dims_mappings = {{1, -1}, {1, -1}, {-1, -1}};
|
||||
expected_dims_mappings = {{{1, -1}, {1, -1}, {1, -1}}, {{1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
forward_input_shapes, input_dims_mappings, expected_dims_mappings);
|
||||
|
||||
// mp case, backward
|
||||
input_dims_mappings = {{1, -1}, {1, -1}, {-1, -1}, {1, -1}};
|
||||
expected_dims_mappings = {{{1, -1}, {1, -1}, {1, -1}, {1, -1}},
|
||||
{{1, -1}, {1, -1}, {1, -1}}};
|
||||
test_moe_combine_spmd(
|
||||
backward_input_shapes, input_dims_mappings, expected_dims_mappings, true);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,188 @@
|
||||
/* Copyright (c) 2024 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using phi::distributed::ArgDistAttr;
|
||||
using phi::distributed::DistMetaTensor;
|
||||
|
||||
void test_moe_gate_dispatch_spmd(
|
||||
const std::vector<std::vector<int64_t>>& input_shapes,
|
||||
const std::vector<std::vector<int64_t>>& input_dims_mappings,
|
||||
const std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>& expected_dims_mappings,
|
||||
int64_t k,
|
||||
int64_t capacity,
|
||||
bool use_pad,
|
||||
bool test_bwd_spmd = false,
|
||||
bool optional = true) {
|
||||
size_t num_inputs = 0;
|
||||
if (test_bwd_spmd) {
|
||||
num_inputs = 5;
|
||||
} else {
|
||||
num_inputs = 3;
|
||||
}
|
||||
|
||||
EXPECT_EQ(input_shapes.size(), num_inputs)
|
||||
<< "The number of input_shapes must be" << num_inputs << ", but got "
|
||||
<< input_shapes.size();
|
||||
EXPECT_EQ(input_dims_mappings.size(), num_inputs)
|
||||
<< "The number of input_dims_mapping must be" << num_inputs
|
||||
<< ", but got " << input_dims_mappings.size();
|
||||
|
||||
std::vector<int64_t> mesh_shape = {4};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"dpmp"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<DistMetaTensor> dist_meta_tensors;
|
||||
for (size_t i = 0; i < num_inputs; ++i) {
|
||||
TensorDistAttr dist_attr = TensorDistAttr();
|
||||
dist_attr.set_process_mesh(process_mesh);
|
||||
|
||||
const std::vector<int64_t>& shape = input_shapes[i];
|
||||
const std::vector<int64_t>& dim_mapping = input_dims_mappings[i];
|
||||
EXPECT_EQ(shape.size(), dim_mapping.size())
|
||||
<< "The size of shape and dim_mapping for input " << i
|
||||
<< " must be the same, but got " << shape.size()
|
||||
<< " != " << dim_mapping.size();
|
||||
|
||||
dist_attr.set_dims_mapping(dim_mapping);
|
||||
dist_attr.set_dynamic_dims(std::vector<bool>(shape.size(), false));
|
||||
|
||||
dist_meta_tensors.push_back(
|
||||
DistMetaTensor(common::make_ddim(shape), dist_attr));
|
||||
}
|
||||
|
||||
phi::distributed::SpmdInfo spmd_info;
|
||||
if (test_bwd_spmd) {
|
||||
spmd_info =
|
||||
phi::distributed::MoEGateDispatchGradInferSpmd(dist_meta_tensors[0],
|
||||
dist_meta_tensors[1],
|
||||
dist_meta_tensors[2],
|
||||
dist_meta_tensors[3],
|
||||
dist_meta_tensors[4],
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
} else {
|
||||
phi::distributed::DistMetaTensor uninitialized_tensor;
|
||||
spmd_info = phi::distributed::MoEGateDispatchInferSpmd(
|
||||
dist_meta_tensors[0],
|
||||
dist_meta_tensors[1],
|
||||
optional ? dist_meta_tensors[2] : uninitialized_tensor,
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < 2; ++i) {
|
||||
std::vector<ArgDistAttr> dist_attrs;
|
||||
std::vector<std::vector<int64_t>> dims_mappings;
|
||||
if (i == 0) {
|
||||
dist_attrs = spmd_info.first;
|
||||
dims_mappings = expected_dims_mappings.first;
|
||||
} else {
|
||||
dist_attrs = spmd_info.second;
|
||||
dims_mappings = expected_dims_mappings.second;
|
||||
}
|
||||
EXPECT_EQ(dist_attrs.size(), dims_mappings.size())
|
||||
<< "The size of dist_attr and expected_dims must be the same, but got "
|
||||
<< dist_attrs.size() << " != " << dims_mappings.size();
|
||||
|
||||
for (size_t j = 0; j < dist_attrs.size(); ++j) {
|
||||
const ArgDistAttr& inferred_attr = dist_attrs[j];
|
||||
const std::vector<int64_t>& expected_dims_mapping = dims_mappings[j];
|
||||
check_dim_mapping(inferred_attr, expected_dims_mapping);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(MoECombineSPMDRule, test_moe_gate_dispatch_spmd) {
|
||||
int64_t s = 1024, h = 512, k = 2, e = 8, capacity = 1024;
|
||||
bool use_pad = true;
|
||||
const std::vector<std::vector<int64_t>>& forward_input_shapes = {
|
||||
{s, h}, {s, e}, {e}};
|
||||
const std::vector<std::vector<int64_t>>& backward_input_shapes = {
|
||||
{s, k}, {k, s}, {s, k}, {e, capacity, h}, {s, k}};
|
||||
|
||||
// replicated case, forward
|
||||
std::vector<std::vector<int64_t>> input_dims_mappings = {
|
||||
{-1, -1}, {-1, -1}, {-1}};
|
||||
std::pair<std::vector<std::vector<int64_t>>,
|
||||
std::vector<std::vector<int64_t>>>
|
||||
expected_dims_mappings = {
|
||||
{{-1, -1}, {-1, -1}, {-1}},
|
||||
{{-1, -1, -1}, {-1, -1}, {-1, -1}, {-1}, {-1, -1}}};
|
||||
test_moe_gate_dispatch_spmd(forward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
|
||||
// replicated case, backward
|
||||
input_dims_mappings = {{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1, -1}, {-1, -1}};
|
||||
expected_dims_mappings = {
|
||||
{{-1, -1}, {-1, -1}, {-1, -1}, {-1, -1, -1}, {-1, -1}},
|
||||
{{-1, -1}, {-1, -1}}};
|
||||
test_moe_gate_dispatch_spmd(backward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad,
|
||||
true);
|
||||
|
||||
// ep case, forward
|
||||
input_dims_mappings = {{0, -1}, {-1, -1}, {-1}};
|
||||
expected_dims_mappings = {{{0, -1}, {0, -1}, {-1}},
|
||||
{{-1, 0, -1}, {0, -1}, {-1, 0}, {-1}, {0, -1}}};
|
||||
test_moe_gate_dispatch_spmd(forward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad);
|
||||
|
||||
// ep case, backward
|
||||
input_dims_mappings = {{0, -1}, {-1, 0}, {0, -1}, {-1, 0, -1}, {0, -1}};
|
||||
expected_dims_mappings = {{{0, -1}, {-1, 0}, {0, -1}, {-1, 0, -1}, {0, -1}},
|
||||
{{0, -1}, {0, -1}}};
|
||||
test_moe_gate_dispatch_spmd(backward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad,
|
||||
true);
|
||||
|
||||
// ep, corr_bias is none case, forward
|
||||
input_dims_mappings = {{0, -1}, {-1, -1}, {-1}};
|
||||
expected_dims_mappings = {{{0, -1}, {0, -1}, {}},
|
||||
{{-1, 0, -1}, {0, -1}, {-1, 0}, {-1}, {0, -1}}};
|
||||
test_moe_gate_dispatch_spmd(forward_input_shapes,
|
||||
input_dims_mappings,
|
||||
expected_dims_mappings,
|
||||
k,
|
||||
capacity,
|
||||
use_pad,
|
||||
false,
|
||||
false);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,54 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
|
||||
|
||||
namespace phi {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(ProcessMesh, Ctor) {
|
||||
std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
int64_t size = shape[0] * shape[1];
|
||||
ProcessMesh process_mesh(shape, process_ids, dim_names);
|
||||
EXPECT_EQ(process_mesh.shape(), shape);
|
||||
EXPECT_EQ(process_mesh.process_ids(), process_ids);
|
||||
EXPECT_EQ(process_mesh.dim_names()[0], "x");
|
||||
EXPECT_EQ(process_mesh.dim_names()[1], "y");
|
||||
EXPECT_EQ(process_mesh.size(), size);
|
||||
EXPECT_EQ(process_mesh.ndim(), static_cast<int64_t>(shape.size()));
|
||||
EXPECT_EQ(process_mesh.dim_size(0), shape[0]);
|
||||
EXPECT_EQ(process_mesh.dim_size(-1), shape[1]);
|
||||
EXPECT_EQ(process_mesh.dim_size("x"), shape[0]);
|
||||
EXPECT_EQ(process_mesh.dim_size("y"), shape[1]);
|
||||
EXPECT_EQ(process_mesh.empty(), false);
|
||||
EXPECT_EQ(process_mesh.contains(0), true);
|
||||
EXPECT_EQ(process_mesh.contains(6), false);
|
||||
std::stringstream sstream;
|
||||
sstream << process_mesh;
|
||||
EXPECT_EQ(sstream.str(), process_mesh.to_string());
|
||||
auto proto = phi::distributed::to_proto(process_mesh);
|
||||
ProcessMesh new_process_mesh = ProcessMesh::from_proto(proto);
|
||||
EXPECT_EQ(process_mesh, new_process_mesh);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,153 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct ReshapeTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> input_dims_mapping;
|
||||
|
||||
// shape attribute
|
||||
std::vector<int64_t> target_shape;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_input_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
};
|
||||
|
||||
TEST(Reshape, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
// test flatten
|
||||
// [4, 6, 8] -> [192]:
|
||||
// [[0], [1], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1]]
|
||||
|
||||
// [4, 6, 8] -> [192]: [[ ], [0], [1]] -> [[ ], [ ], [ ]], [[ ]]
|
||||
|
||||
// [4, 6, 8] -> [192]:
|
||||
// [[0, 1], [ ], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1]]
|
||||
|
||||
// [2, 12, 8] -> [192]:
|
||||
// [[0], [1], [ ]] -> [[0], [ ], [ ]], [[0]]
|
||||
|
||||
// test split
|
||||
|
||||
// [128] -> [4, 6, 8]:
|
||||
// [[0, 1]] -> [[0, 1]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [128] -> [6, 4, 8]:
|
||||
// [[0, 1]] -> [[ ], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [2, 12, 8]
|
||||
// [[0], [1], [ ]] -> [[0], [ ], [ ]], [[0], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [2, 12, 8]
|
||||
// [[0, 1], [ ], [ ]] -> [[ ], [ ], [ ]], [[ ], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [12, 2, 8]:
|
||||
// [[0], [1], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [12, 2, 8]:
|
||||
// [[0, 1], [ ], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[0], [1], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[ ], [0], [1]] -> [[ ], [ ], [ ]], [[ ], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[0], [ ], [1]] -> [[0], [ ], [ ]], [[0], [ ], [ ]]
|
||||
|
||||
// [4, 6, 8] -> [8, 6, 4]:
|
||||
// [[0, 1], [ ], [ ]] -> [[0, 1], [ ], [ ]], [[0, 1], [ ], [ ]]
|
||||
|
||||
std::vector<ReshapeTestCase> test_cases = {
|
||||
// input_shape, input_dims_mapping, target_shape,
|
||||
// expected_input_dims_mapping, expected_output_dims_mapping
|
||||
|
||||
// test flatten
|
||||
{{4, 6, 8}, {{0}, {1}, {}}, {192}, {{0, 1}, {}, {}}, {{0, 1}}},
|
||||
{{4, 6, 8}, {{}, {0}, {1}}, {192}, {{}, {}, {}}, {{}}},
|
||||
{{4, 6, 8}, {{0, 1}, {}, {}}, {192}, {{0, 1}, {}, {}}, {{0, 1}}},
|
||||
{{2, 12, 8}, {{0}, {1}, {}}, {192}, {{0}, {}, {}}, {{0}}},
|
||||
|
||||
// test split
|
||||
{{192}, {{0, 1}}, {4, 6, 8}, {{0, 1}}, {{0, 1}, {}, {}}},
|
||||
{{192}, {{0, 1}}, {6, 4, 8}, {{}}, {{}, {}, {}}},
|
||||
|
||||
// test combination
|
||||
{{4, 6, 8}, {{0}, {1}, {}}, {2, 12, 8}, {{0}, {}, {}}, {{0}, {}, {}}},
|
||||
{{4, 6, 8}, {{0, 1}, {}, {}}, {2, 12, 8}, {{}, {}, {}}, {{}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0}, {1}, {}},
|
||||
{12, 2, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{12, 2, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0}, {1}, {}},
|
||||
{8, 6, 4},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8}, {{}, {0}, {1}}, {8, 6, 4}, {{}, {}, {}}, {{}, {}, {}}},
|
||||
{{4, 6, 8}, {{0}, {}, {1}}, {8, 6, 4}, {{0}, {}, {}}, {{0}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{0, 1}, {}, {}},
|
||||
{8, 6, 4},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
{{4, 6, 8}, {{}, {}, {0, 1}}, {24, 2, 4}, {{}, {}, {}}, {{}, {}, {}}},
|
||||
{{4, 6, 8},
|
||||
{{}, {}, {0, 1}},
|
||||
{24, 4, 2},
|
||||
{{}, {}, {0, 1}},
|
||||
{{}, {0, 1}, {}}},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.input_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::ReshapeInferSpmd(x, tc.target_shape);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_input_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
}
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,243 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct SoftmaxTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> input_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_input_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
};
|
||||
|
||||
struct SoftmaxGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> out_shape;
|
||||
std::vector<std::vector<int64_t>> out_dims_mapping;
|
||||
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// axis attribute
|
||||
int axis;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
};
|
||||
|
||||
TEST(SoftmaxInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<SoftmaxTestCase> test_cases = {
|
||||
// shape = [32, 48, 128], axis = 0
|
||||
// [[0,1],[2],[]] -> [[],[2],[]], [[],[2],[]]
|
||||
{{32, 48, 128}, {{0, 1}, {2}, {}}, 0, {{}, {2}, {}}, {{}, {2}, {}}},
|
||||
{{32, 48, 128}, {{0, 1}, {2}, {}}, -3, {{}, {2}, {}}, {{}, {2}, {}}},
|
||||
|
||||
// shape = [32, 48, 128], axis = 1
|
||||
// [[0,1],[2],[]] -> [[0, 1],[],[]], [[0, 1],[],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
1,
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}}};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.input_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::SoftmaxInferSpmd(x, tc.axis);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_input_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(SoftmaxGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<SoftmaxGradTestCase> test_cases = {
|
||||
// out_shape = [32, 48, 128], out_grad_shape = [32, 48, 128], axis = 0
|
||||
// [[0,1],[2],[]], [[0,1],[2],[]] -> [[],[2],[]], [[],[2],[]], [[],[2],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
0,
|
||||
{{}, {2}, {}},
|
||||
{{}, {2}, {}},
|
||||
{{}, {2}, {}}},
|
||||
// axis = 0
|
||||
// [[0,1],[2],[]], [[0],[1,2],[]] -> [[],[1,2],[]], [[],[1, 2],[]],
|
||||
// [[],[1,2],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
{32, 48, 128},
|
||||
{{0}, {1, 2}, {}},
|
||||
0,
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}}},
|
||||
// axis = 1
|
||||
// [[0,1],[2],[]], [[2],[0,1],[]] -> [[0,1,2],[],[]], [[0, 1, 2],[],[]],
|
||||
// [[0, 1, 2],[],[]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {2}, {}},
|
||||
{32, 48, 128},
|
||||
{{2}, {0, 1}, {}},
|
||||
1,
|
||||
{{0, 1, 2}, {}, {}},
|
||||
{{0, 1, 2}, {}, {}},
|
||||
{{0, 1, 2}, {}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1],[]], [[],[0,1],[]] -> [[],[0,1],[]], [[],[0,1],[]],
|
||||
// [[],[0,1],[]]
|
||||
{{32, 48, 128},
|
||||
{{0}, {1}, {}},
|
||||
{32, 48, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1],[]], [[0,1],[],[]] -> [[0,1],[],[]], [[0, 1],[],[]],
|
||||
// [[0,1],[],[]]
|
||||
{{32, 48, 128},
|
||||
{{0}, {1}, {}},
|
||||
{32, 48, 128},
|
||||
{{0, 1}, {}, {}},
|
||||
2,
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1,2],[]], [[],[0,1],[]] -> [[0],[1,2],[]], [[0],[1,2],[]],
|
||||
// [[0],[1,2],[]]
|
||||
{{32, 48, 128},
|
||||
{{0}, {1, 2}, {}},
|
||||
{32, 48, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}}},
|
||||
// axis = 2
|
||||
// [[0],[1,2],[]], [[],[0,1],[]] -> [[0],[1,2],[]], [[0],[1,2],[]],
|
||||
// [[0],[1,2],[]]
|
||||
{{2, 4, 128},
|
||||
{{0}, {1, 2}, {}},
|
||||
{2, 4, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}},
|
||||
{{0}, {1, 2}, {}}},
|
||||
// axis = 2
|
||||
// [[],[1,2],[]], [[],[0,1],[]] -> [[],[1,2],[]], [[],[1,2],[]],
|
||||
// [[],[1,2],[]]
|
||||
{{2, 4, 128},
|
||||
{{}, {1, 2}, {}},
|
||||
{2, 4, 128},
|
||||
{{}, {0, 1}, {}},
|
||||
2,
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}},
|
||||
{{}, {1, 2}, {}}},
|
||||
// axis = 1
|
||||
// [[0,1],[],[]], [[],[],[2]] -> [[0,1],[],[2]], [[0,1],[],[2]],
|
||||
// [[0,1],[],[2]]
|
||||
{{32, 48, 128},
|
||||
{{0, 1}, {}, {}},
|
||||
{32, 48, 128},
|
||||
{{}, {}, {2}},
|
||||
1,
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}}},
|
||||
// Note: just for pass coverage ci: axis = 2
|
||||
// [[0],[0,1],[]], [[],[],[]] -> [[],[0,1],[]], [[],[0,1],[]],
|
||||
// [[],[0,1],[]]
|
||||
{{2, 4, 128},
|
||||
{{0}, {0, 1}, {}},
|
||||
{2, 4, 128},
|
||||
{{}, {}, {}},
|
||||
2,
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}}}};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr out_dist_attr = TensorDistAttr();
|
||||
out_dist_attr.set_process_mesh(process_mesh);
|
||||
out_dist_attr.set_dims_mapping(tc.out_dims_mapping);
|
||||
out_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.out_shape), out_dist_attr);
|
||||
TensorDistAttr out_grad_attr = TensorDistAttr();
|
||||
out_grad_attr.set_process_mesh(process_mesh);
|
||||
out_grad_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::SoftmaxGradInferSpmd(out, out_grad, tc.axis);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
// [[0,1],[2]] [[2],[]]
|
||||
@@ -0,0 +1,143 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(SoftmaxGradInferSpmd, Ctor) {
|
||||
// Sharding along axes besides softmax axis.
|
||||
std::vector<int64_t> x_shape = {36, 48};
|
||||
std::vector<int64_t> out_grad_shape = {36, 48};
|
||||
|
||||
std::vector<int64_t> mesh_shape = {2, 3};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
TensorDistAttr out_grad_dist_attr = TensorDistAttr();
|
||||
out_grad_dist_attr.set_process_mesh(process_mesh);
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({1, -1}));
|
||||
out_grad_dist_attr.set_dynamic_dims(std::vector<bool>({false, false}));
|
||||
|
||||
phi::distributed::DistMetaTensor x(phi::make_ddim(x_shape), x_dist_attr);
|
||||
phi::distributed::DistMetaTensor out_grad(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
int axis = 1;
|
||||
|
||||
auto spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]), std::vector<int64_t>({1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]), std::vector<int64_t>({1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on other axes." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// Sharding along softmax axis.
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1}));
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, 1}));
|
||||
x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
axis = 1;
|
||||
|
||||
spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({-1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on softmax axis." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// Sharding on multi axes.
|
||||
x_shape = {10, 36, 48, 24};
|
||||
out_grad_shape = {10, 36, 48, 24};
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({0, 1, -1, -1}));
|
||||
x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
axis = 1;
|
||||
|
||||
spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({0, -1, -1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, -1, -1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, -1, -1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on multi axes." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
|
||||
// Sharding on multi axes.
|
||||
x_shape = {10, 36, 48, 24};
|
||||
out_grad_shape = {10, 36, 48, 24};
|
||||
x_dist_attr.set_dims_mapping(std::vector<int64_t>({0, -1, -1, -1}));
|
||||
out_grad_dist_attr.set_dims_mapping(std::vector<int64_t>({-1, -1, 1, -1}));
|
||||
x = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape), x_dist_attr);
|
||||
out_grad = phi::distributed::DistMetaTensor(phi::make_ddim(x_shape),
|
||||
out_grad_dist_attr);
|
||||
axis = 1;
|
||||
|
||||
spmdinfo = SoftmaxGradInferSpmd(x, out_grad, axis);
|
||||
|
||||
EXPECT_EQ(spmdinfo.first.size(), 2UL);
|
||||
EXPECT_EQ(spmdinfo.second.size(), 1UL);
|
||||
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[0]),
|
||||
std::vector<int64_t>({0, -1, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.first[1]),
|
||||
std::vector<int64_t>({0, -1, 1, -1}));
|
||||
EXPECT_EQ(get_dims_mapping(spmdinfo.second[0]),
|
||||
std::vector<int64_t>({0, -1, 1, -1}));
|
||||
EXPECT_DOUBLE_EQ(
|
||||
PADDLE_GET_CONST(TensorDistAttr, spmdinfo.second[0]).is_partial(), false);
|
||||
VLOG(4) << "Test SoftmaxGradInferSpmd sharding on multi axes." << std::endl
|
||||
<< std::endl
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,118 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
const std::vector<int64_t>& get_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.dims_mapping();
|
||||
}
|
||||
|
||||
bool is_partial(const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.is_partial();
|
||||
}
|
||||
|
||||
const std::set<int64_t> get_partial_dims(
|
||||
const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.partial_dims();
|
||||
}
|
||||
|
||||
const std::vector<std::vector<int64_t>>& get_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr));
|
||||
const auto& tensor_attr =
|
||||
PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
|
||||
return tensor_attr.multi_dims_mapping();
|
||||
}
|
||||
|
||||
void check_dim_mapping(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<int64_t>& dim_mapping,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(get_dims_mapping(dist_attr), dim_mapping) << line;
|
||||
}
|
||||
|
||||
void check_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<std::vector<int64_t>>& dim_mapping,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(get_multi_dims_mapping(dist_attr), dim_mapping) << line;
|
||||
}
|
||||
|
||||
void check_empty_dist_attr(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr),
|
||||
phi::distributed::TensorDistAttr());
|
||||
}
|
||||
|
||||
void check_partial_dims(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::set<int64_t>& dims,
|
||||
const std::string& line) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr))
|
||||
<< line;
|
||||
EXPECT_EQ(get_partial_dims(dist_attr), dims) << line;
|
||||
}
|
||||
|
||||
void clean_partial_status(phi::distributed::ArgDistAttr* dist_attr) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
|
||||
auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
|
||||
tensor_attr.clean_partial_status();
|
||||
}
|
||||
|
||||
void clean_partial_dims(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
|
||||
auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
|
||||
tensor_attr.clean_partial_dims(dims);
|
||||
}
|
||||
|
||||
void set_partial_status(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims) {
|
||||
EXPECT_TRUE(
|
||||
paddle::holds_alternative<phi::distributed::TensorDistAttr>(*dist_attr));
|
||||
auto& tensor_attr = PADDLE_GET(phi::distributed::TensorDistAttr, *dist_attr);
|
||||
tensor_attr.set_partial_status(dims);
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,74 @@
|
||||
/* Copyright (c) 2022 PaddlePaddle 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. */
|
||||
#pragma once
|
||||
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
#include "glog/logging.h"
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
|
||||
#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
|
||||
#include "paddle/phi/core/distributed/type_defs.h"
|
||||
#include "paddle/phi/infermeta/spmd_rules/replicated.h"
|
||||
#include "paddle/phi/infermeta/spmd_rules/rules.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
using phi::distributed::ProcessMesh;
|
||||
using phi::distributed::TensorDistAttr;
|
||||
|
||||
const std::vector<int64_t>& get_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
bool is_partial(const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
const std::set<int64_t> get_partial_dims(
|
||||
const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
const std::vector<std::vector<int64_t>>& get_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr);
|
||||
|
||||
void check_dim_mapping(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<int64_t>& dim_mapping,
|
||||
const std::string& line = "");
|
||||
|
||||
void check_multi_dims_mapping(
|
||||
const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::vector<std::vector<int64_t>>& dim_mapping,
|
||||
const std::string& line = "");
|
||||
|
||||
void check_empty_dist_attr(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::string& line = "");
|
||||
|
||||
void check_partial_dims(const phi::distributed::ArgDistAttr& dist_attr,
|
||||
const std::set<int64_t>& dims,
|
||||
const std::string& line = "");
|
||||
|
||||
void clean_partial_status(phi::distributed::ArgDistAttr* dist_attr);
|
||||
|
||||
void clean_partial_dims(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims);
|
||||
|
||||
void set_partial_status(phi::distributed::ArgDistAttr* dist_attr,
|
||||
std::vector<int64_t> dims);
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,215 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 "paddle/phi/infermeta/spmd_rules/tile.h"
|
||||
#include "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct TileTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
// repeat_times attribute
|
||||
phi::IntArray repeat_times;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_dims_mapping;
|
||||
};
|
||||
|
||||
struct TileGradTestCase {
|
||||
// input
|
||||
std::vector<int64_t> x_shape;
|
||||
std::vector<std::vector<int64_t>> x_dims_mapping;
|
||||
|
||||
std::vector<int64_t> out_grad_shape;
|
||||
std::vector<std::vector<int64_t>> out_grad_dims_mapping;
|
||||
|
||||
// repeat_times attribute
|
||||
phi::IntArray repeat_times;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_x_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_out_grad_dims_mapping;
|
||||
|
||||
std::vector<std::vector<int64_t>> expected_x_grad_dims_mapping;
|
||||
|
||||
std::set<int64_t> partial_dims;
|
||||
};
|
||||
|
||||
TEST(TileInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<TileTestCase> test_cases = {
|
||||
// shape = [8, 16, 24], repeat_times = {2, 2, 1, 1}
|
||||
// [[0],[],[1,2]] -> [[],[],[1,2]], [[],[],[],[1,2]]
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0}, {}, {1, 2}},
|
||||
phi::IntArray({2, 2, 1, 1}),
|
||||
{{}, {}, {1, 2}},
|
||||
{{}, {}, {}, {1, 2}},
|
||||
},
|
||||
|
||||
// shape = [8, 16, 24], repeat_times = {1, 2}
|
||||
// [[0,1],[],[2]] -> [[0,1],[],[]], [[0,1],[],[]]
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
phi::IntArray({1, 2}),
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
},
|
||||
|
||||
// shape = [8, 16, 24], repeat_times = {}
|
||||
// [[0,1],[],[2]] -> [[0,1],[],[2]], [[0,1],[],[2]]
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
phi::IntArray({}),
|
||||
{{0, 1}, {}, {2}},
|
||||
{{0, 1}, {}, {2}},
|
||||
},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::TileInferSpmdDynamic(x, tc.repeat_times);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_out_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TileGradInferSpmd, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5, 6, 7};
|
||||
std::vector<std::string> dim_names = {"x", "y", "z"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<TileGradTestCase> test_cases = {
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [2, 16, 16, 24], repeat_times =
|
||||
// {2, 2, 1, 1}
|
||||
// [[0],[],[1,2]], [[],[],[],[1,2]] -> [[],[],[1,2]], [[],[],[],[1,2]],
|
||||
// [[],[],[1,2]], partial on {}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0}, {}, {1, 2}},
|
||||
{2, 16, 16, 24},
|
||||
{{}, {}, {}, {1, 2}},
|
||||
phi::IntArray({2, 2, 1, 1}),
|
||||
{{}, {}, {1, 2}},
|
||||
{{}, {}, {}, {1, 2}},
|
||||
{{}, {}, {1, 2}},
|
||||
{},
|
||||
},
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [8, 16, 48], repeat_times = {1,
|
||||
// 2}
|
||||
// [[0,1],[],[2]], [[0,1],[],[2]] -> [[0,1],[],[]], [[0,1],[],[]]],
|
||||
// [[0,1],[],[]], partial on {}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
{8, 16, 48},
|
||||
{{0, 1}, {}, {2}},
|
||||
phi::IntArray({1, 2}),
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{{0, 1}, {}, {}},
|
||||
{},
|
||||
},
|
||||
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [8, 16, 24], repeat_times = {}
|
||||
// [[0,1],[],[2]], [[0],[1],[2]] -> [[0],[1],[2]], [[0],[1],[2]],
|
||||
// [[0],[1],[2]], partial on {}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0, 1}, {}, {2}},
|
||||
{8, 16, 24},
|
||||
{{0}, {1}, {2}},
|
||||
phi::IntArray({}),
|
||||
{{0}, {1}, {2}},
|
||||
{{0}, {1}, {2}},
|
||||
{{0}, {1}, {2}},
|
||||
{},
|
||||
},
|
||||
|
||||
// x_shape = [8, 16, 24], out_grad_shape = [8, 16, 16, 24], repeat_times =
|
||||
// {8, 2, 1, 1}
|
||||
// [[0],[],[]], [[1,2],[],[],[]] -> [[],[],[]], [[1,2],[],[],[]],
|
||||
// [[],[],[]], partial on {1,2}
|
||||
{
|
||||
{8, 16, 24},
|
||||
{{0}, {}, {}},
|
||||
{8, 16, 16, 24},
|
||||
{{1, 2}, {}, {}, {}},
|
||||
phi::IntArray({8, 2, 1, 1}),
|
||||
{{}, {}, {}},
|
||||
{{1, 2}, {}, {}, {}},
|
||||
{{}, {}, {}},
|
||||
{1, 2},
|
||||
},
|
||||
};
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr x_dist_attr = TensorDistAttr();
|
||||
x_dist_attr.set_process_mesh(process_mesh);
|
||||
x_dist_attr.set_dims_mapping(tc.x_dims_mapping);
|
||||
x_dist_attr.set_dynamic_dims(std::vector<bool>(tc.x_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.x_shape), x_dist_attr);
|
||||
TensorDistAttr out_grad_attr = TensorDistAttr();
|
||||
out_grad_attr.set_process_mesh(process_mesh);
|
||||
out_grad_attr.set_dims_mapping(tc.out_grad_dims_mapping);
|
||||
out_grad_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.out_grad_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor out_grad =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(tc.out_grad_shape),
|
||||
out_grad_attr);
|
||||
|
||||
// test backward
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
phi::distributed::TileGradInferSpmdDynamic(
|
||||
x, out_grad, tc.repeat_times);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(backward_spmd_info.first[0],
|
||||
tc.expected_x_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.first[1],
|
||||
tc.expected_out_grad_dims_mapping);
|
||||
check_multi_dims_mapping(backward_spmd_info.second[0],
|
||||
tc.expected_x_grad_dims_mapping);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,63 @@
|
||||
/* Copyright (c) 2023 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
TEST(Tile, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
std::vector<int64_t> shape = {6, 8, 10};
|
||||
std::vector<int64_t> dims_mapping = {0, -1, 1};
|
||||
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims({false, false, false});
|
||||
phi::distributed::DistMetaTensor x =
|
||||
phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
|
||||
std::vector<int64_t> repeat_times = {2, 2, 1, 1};
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::TileInferSpmd(x, repeat_times);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_dim_mapping(forward_spmd_info.first[0], {-1, -1, 1});
|
||||
check_dim_mapping(forward_spmd_info.second[0], {-1, -1, -1, 1});
|
||||
check_partial_dims(forward_spmd_info.second[0], {});
|
||||
|
||||
// test backward
|
||||
auto out_grad_dist_attr =
|
||||
PADDLE_GET_CONST(TensorDistAttr, forward_spmd_info.second[0]);
|
||||
out_grad_dist_attr.set_dims_mapping({0, -1, -1, 1});
|
||||
phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim({2, 12, 8, 10}), out_grad_dist_attr);
|
||||
phi::distributed::SpmdInfo backward_spmd_info =
|
||||
TileGradInferSpmd(x, out_grad, repeat_times);
|
||||
EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(2));
|
||||
EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_dim_mapping(backward_spmd_info.first[0], {-1, -1, 1});
|
||||
check_dim_mapping(backward_spmd_info.first[1], {0, -1, -1, 1});
|
||||
check_dim_mapping(backward_spmd_info.second[0], {-1, -1, 1});
|
||||
check_partial_dims(backward_spmd_info.second[0], {0});
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,89 @@
|
||||
/* Copyright (c) 2025 PaddlePaddle 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 "test/cpp/auto_parallel/spmd_rule_test_util.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace distributed {
|
||||
namespace auto_parallel {
|
||||
|
||||
struct TransposeTestCase {
|
||||
// input
|
||||
std::vector<int64_t> input_shape;
|
||||
std::vector<std::vector<int64_t>> input_dims_mapping;
|
||||
|
||||
// shape attribute
|
||||
std::vector<int> perm;
|
||||
|
||||
// output
|
||||
std::vector<std::vector<int64_t>> expected_input_dims_mapping;
|
||||
std::vector<std::vector<int64_t>> expected_output_dims_mapping;
|
||||
};
|
||||
|
||||
TEST(Transpose, Ctor) {
|
||||
std::vector<int64_t> mesh_shape = {2, 2};
|
||||
std::vector<int64_t> process_ids = {0, 1, 2, 3};
|
||||
std::vector<std::string> dim_names = {"x", "y"};
|
||||
ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
|
||||
|
||||
//
|
||||
std::vector<TransposeTestCase> test_cases = {
|
||||
// input_shape, input_dims_mapping, perm,
|
||||
// expected_input_dims_mapping, expected_output_dims_mapping
|
||||
|
||||
{{64, 48, 36, 24},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{1, 0, 2, 3},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{}, {0, 1}, {}, {}}},
|
||||
{{64, 48, 36, 24},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{0, 1, 2, 3},
|
||||
{{0, 1}, {}, {}, {}},
|
||||
{{0, 1}, {}, {}, {}}},
|
||||
{{64, 48, 36, 24},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{0, 2, 3, 1},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{{}, {0, 1}, {}, {}}},
|
||||
{{64, 48, 36, 24},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{-1, 0, -2, 1},
|
||||
{{}, {}, {0, 1}, {}},
|
||||
{{}, {}, {0, 1}, {}}},
|
||||
};
|
||||
|
||||
for (const auto& tc : test_cases) {
|
||||
TensorDistAttr t_dist_attr = TensorDistAttr();
|
||||
t_dist_attr.set_process_mesh(process_mesh);
|
||||
t_dist_attr.set_dims_mapping(tc.input_dims_mapping);
|
||||
t_dist_attr.set_dynamic_dims(
|
||||
std::vector<bool>(tc.input_shape.size(), false));
|
||||
phi::distributed::DistMetaTensor x = phi::distributed::DistMetaTensor(
|
||||
common::make_ddim(tc.input_shape), t_dist_attr);
|
||||
|
||||
// test forward
|
||||
phi::distributed::SpmdInfo forward_spmd_info =
|
||||
phi::distributed::TransposeInferSpmd(x, tc.perm);
|
||||
EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
|
||||
EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
|
||||
check_multi_dims_mapping(forward_spmd_info.first[0],
|
||||
tc.expected_input_dims_mapping);
|
||||
check_multi_dims_mapping(forward_spmd_info.second[0],
|
||||
tc.expected_output_dims_mapping);
|
||||
}
|
||||
}
|
||||
} // namespace auto_parallel
|
||||
} // namespace distributed
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,9 @@
|
||||
add_subdirectory(adt)
|
||||
add_subdirectory(ast_gen_ius)
|
||||
add_subdirectory(common)
|
||||
add_subdirectory(hlir)
|
||||
add_subdirectory(ir)
|
||||
add_subdirectory(lang)
|
||||
add_subdirectory(optim)
|
||||
add_subdirectory(runtime)
|
||||
add_subdirectory(utils)
|
||||
@@ -0,0 +1,4 @@
|
||||
cinn_cc_test(equation_value_match_trait_test SRCS
|
||||
equation_value_match_trait_test.cc DEPS gtest glog)
|
||||
|
||||
cinn_cc_test(tree_test SRCS tree_test.cc DEPS gtest glog)
|
||||
@@ -0,0 +1,143 @@
|
||||
// Copyright (c) 2023 PaddlePaddle 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 "paddle/cinn/adt/equation_value_match_trait.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/cinn/adt/equation_value.h"
|
||||
#include "paddle/cinn/adt/match.h"
|
||||
|
||||
namespace cinn::adt::test {
|
||||
|
||||
TEST(Match, Union) {
|
||||
using Pattern = Union<IndexUnDotValue<Value, List<DimExpr>>,
|
||||
IndexDotValue<Value, List<DimExpr>>>;
|
||||
{
|
||||
Value expr = IndexUnDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
|
||||
bool ret = cinn::adt::Match<Pattern>(expr);
|
||||
ASSERT_TRUE(ret);
|
||||
}
|
||||
{
|
||||
Value expr = IndexDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
|
||||
bool ret = cinn::adt::Match<Pattern>(expr);
|
||||
ASSERT_TRUE(ret);
|
||||
}
|
||||
{
|
||||
Value expr = List<Value>{Value{Ok()}, Value{Ok()}, Value{Ok()}};
|
||||
|
||||
bool ret = cinn::adt::Match<Pattern>(expr);
|
||||
ASSERT_FALSE(ret);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Match, index_undot) {
|
||||
Value expr = IndexUnDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
|
||||
bool ret = cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(expr);
|
||||
ASSERT_TRUE(ret);
|
||||
}
|
||||
|
||||
TEST(Match, index_dot) {
|
||||
Value expr = IndexDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
|
||||
bool ret = cinn::adt::Match<IndexDotValue<Value, List<DimExpr>>>(expr);
|
||||
ASSERT_TRUE(ret);
|
||||
}
|
||||
|
||||
TEST(Match, list) {
|
||||
Value expr = List<Value>{Value{Ok()}, Value{Ok()}, Value{Ok()}};
|
||||
|
||||
bool ret = cinn::adt::Match<List<Value>>(expr);
|
||||
ASSERT_TRUE(ret);
|
||||
}
|
||||
|
||||
TEST(Match, list_get_item) {
|
||||
Value list = List<Value>{Value{Ok()}, Value{Ok()}, Value{Ok()}};
|
||||
Value expr = ListGetItem<Value, DimExpr>{list, DimExpr{std::int64_t(1)}};
|
||||
|
||||
bool ret = cinn::adt::Match<ListGetItem<Value, std::int64_t>>(expr);
|
||||
ASSERT_TRUE(ret);
|
||||
}
|
||||
|
||||
TEST(Match, list_get_item_index_undot) {
|
||||
Value undot1 = IndexUnDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(undot1)));
|
||||
|
||||
Value expr = ListGetItem<Value, DimExpr>{undot1, DimExpr{std::int64_t(1)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
|
||||
expr)));
|
||||
}
|
||||
|
||||
// List<ListGetItem<IndexUnDotValue<Value>, std::int64_t>>
|
||||
TEST(Match, list_list_get_item_index_undot) {
|
||||
Value undot = IndexUnDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
ASSERT_TRUE((cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(undot)));
|
||||
Value expr1 = ListGetItem<Value, DimExpr>{undot, DimExpr{std::int64_t(0)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
|
||||
expr1)));
|
||||
Value expr2 = ListGetItem<Value, DimExpr>{undot, DimExpr{std::int64_t(1)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
|
||||
expr2)));
|
||||
Value list = List<Value>{expr1, expr2};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<List<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>>(
|
||||
list)));
|
||||
}
|
||||
|
||||
// IndexDotValue<List<ListGetItem<IndexUnDotValue<Value>, std::int64_t>>>
|
||||
TEST(Match, index_dot_list_list_get_item_index_undot) {
|
||||
Value undot1 = IndexUnDotValue<Value, List<DimExpr>>{
|
||||
Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(undot1)));
|
||||
Value expr1 = ListGetItem<Value, DimExpr>{undot1, DimExpr{std::int64_t(0)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
|
||||
expr1)));
|
||||
Value expr2 = ListGetItem<Value, DimExpr>{undot1, DimExpr{std::int64_t(1)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
|
||||
expr2)));
|
||||
Value list = List<Value>{expr1, expr2};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<List<
|
||||
ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>>(
|
||||
list)));
|
||||
Value dot =
|
||||
IndexDotValue<Value, List<DimExpr>>{list, List<DimExpr>{std::int64_t(1)}};
|
||||
ASSERT_TRUE(
|
||||
(cinn::adt::Match<
|
||||
IndexDotValue<List<ListGetItem<IndexUnDotValue<Value, List<DimExpr>>,
|
||||
std::int64_t>>,
|
||||
List<DimExpr>>>(dot)));
|
||||
}
|
||||
|
||||
} // namespace cinn::adt::test
|
||||
@@ -0,0 +1,199 @@
|
||||
// Copyright (c) 2023 PaddlePaddle 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 "paddle/cinn/adt/tree.h"
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace cinn::adt {
|
||||
|
||||
namespace test {
|
||||
|
||||
using IntTreeLeafT = std::vector<int>;
|
||||
using IntTreeInnerDataT = std::vector<int>;
|
||||
using IntVecTree = Tree<TreeInner<IntTreeInnerDataT>::Node, IntTreeLeafT>;
|
||||
using IntTreeInnerT = TreeInner<IntTreeInnerDataT>::template Node<IntVecTree>;
|
||||
|
||||
} // namespace test
|
||||
|
||||
template <>
|
||||
struct TreeMerger<test::IntVecTree> {
|
||||
using tree_type = test::IntVecTree;
|
||||
using inner_type = typename TreeTrait<test::IntVecTree>::inner_type;
|
||||
using leaf_type = typename TreeTrait<test::IntVecTree>::leaf_type;
|
||||
using inner_data_type = typename inner_type::value_type;
|
||||
|
||||
inner_data_type GetInnerDataForLeaf(const leaf_type& leaf) const {
|
||||
return leaf;
|
||||
}
|
||||
|
||||
inner_type MakeInnerNode(const inner_data_type& inner_data,
|
||||
const List<test::IntVecTree>& children) const {
|
||||
return inner_type{inner_data, children};
|
||||
}
|
||||
|
||||
using MergeResult = std::tuple<tCommon<inner_data_type>,
|
||||
tLhsRemainder<inner_data_type>,
|
||||
tRhsRemainder<inner_data_type>>;
|
||||
|
||||
MergeResult MergeInnerValue(const inner_data_type& lhs,
|
||||
const inner_data_type& rhs) const {
|
||||
inner_data_type common{};
|
||||
inner_data_type lhs_remainder{};
|
||||
inner_data_type rhs_remainder{};
|
||||
int min_size = std::min(lhs.size(), rhs.size());
|
||||
int idx = 0;
|
||||
for (; idx < min_size; ++idx) {
|
||||
if (lhs.at(idx) == rhs.at(idx)) {
|
||||
common.emplace_back(lhs.at(idx));
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
for (int lhs_idx = idx; lhs_idx < lhs.size(); ++lhs_idx) {
|
||||
lhs_remainder.emplace_back(lhs.at(lhs_idx));
|
||||
}
|
||||
for (int rhs_idx = idx; rhs_idx < rhs.size(); ++rhs_idx) {
|
||||
rhs_remainder.emplace_back(rhs.at(rhs_idx));
|
||||
}
|
||||
return MergeResult{common, lhs_remainder, rhs_remainder};
|
||||
}
|
||||
};
|
||||
|
||||
namespace test {
|
||||
|
||||
TEST(IntVecTree, naive) {
|
||||
List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3}, IntTreeLeafT{4, 5, 6}};
|
||||
TreeMerger<test::IntVecTree> tree_merger{};
|
||||
List<IntVecTree> ret = MakeMergedTrees(tree_merger, leaves);
|
||||
ASSERT_EQ(ret->size(), 2);
|
||||
|
||||
ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
|
||||
const auto& [inner_data0, children0] =
|
||||
ret->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
|
||||
ASSERT_TRUE((children0->size() == 1));
|
||||
ASSERT_TRUE((children0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children0->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
|
||||
|
||||
ASSERT_TRUE(ret->at(1).Has<IntTreeInnerT>());
|
||||
const auto& [inner_data1, children1] =
|
||||
ret->at(1).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data1 == IntTreeLeafT{4, 5, 6}));
|
||||
ASSERT_TRUE((children1->size() == 1));
|
||||
ASSERT_TRUE((children1->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children1->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{4, 5, 6}));
|
||||
}
|
||||
|
||||
TEST(IntVecTree, left_equal_right) {
|
||||
List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3}, IntTreeLeafT{1, 2, 3}};
|
||||
List<IntVecTree> ret =
|
||||
MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
|
||||
ASSERT_EQ(ret->size(), 1);
|
||||
|
||||
ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
|
||||
const auto& [inner_data0, children0] =
|
||||
ret->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
|
||||
ASSERT_TRUE((children0->size() == 2));
|
||||
ASSERT_TRUE((children0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children0->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
|
||||
ASSERT_TRUE((children0->at(1).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children0->at(1).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
|
||||
}
|
||||
|
||||
TEST(IntVecTree, left_gt_right) {
|
||||
List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3, 4, 5}, IntTreeLeafT{1, 2, 3}};
|
||||
List<IntVecTree> ret =
|
||||
MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
|
||||
ASSERT_EQ(ret->size(), 1);
|
||||
|
||||
ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
|
||||
const auto& [inner_data0, children0] =
|
||||
ret->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
|
||||
ASSERT_TRUE((children0->size() == 2));
|
||||
|
||||
ASSERT_TRUE((children0->at(0).Has<IntTreeInnerT>()));
|
||||
const auto& [inner_data_left0, children_left0] =
|
||||
children0->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data_left0 == IntTreeLeafT{4, 5}));
|
||||
ASSERT_TRUE((children_left0->size() == 1));
|
||||
ASSERT_TRUE((children_left0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children_left0->at(0).Get<IntTreeLeafT>() ==
|
||||
IntTreeLeafT{1, 2, 3, 4, 5}));
|
||||
|
||||
ASSERT_TRUE((children0->at(1).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children0->at(1).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
|
||||
}
|
||||
|
||||
TEST(IntVecTree, left_lt_right) {
|
||||
List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3}, IntTreeLeafT{1, 2, 3, 4, 5}};
|
||||
List<IntVecTree> ret =
|
||||
MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
|
||||
ASSERT_EQ(ret->size(), 1);
|
||||
|
||||
ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
|
||||
const auto& [inner_data0, children0] =
|
||||
ret->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
|
||||
ASSERT_TRUE((children0->size() == 2));
|
||||
|
||||
ASSERT_TRUE((children0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children0->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
|
||||
|
||||
ASSERT_TRUE((children0->at(1).Has<IntTreeInnerT>()));
|
||||
const auto& [inner_data_right0, children_right0] =
|
||||
children0->at(1).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data_right0 == IntTreeLeafT{4, 5}));
|
||||
ASSERT_TRUE((children_right0->size() == 1));
|
||||
ASSERT_TRUE((children_right0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children_right0->at(0).Get<IntTreeLeafT>() ==
|
||||
IntTreeLeafT{1, 2, 3, 4, 5}));
|
||||
}
|
||||
|
||||
TEST(IntVecTree, left_ne_right) {
|
||||
List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3, 4, 5},
|
||||
IntTreeLeafT{1, 2, 3, 6, 7}};
|
||||
List<IntVecTree> ret =
|
||||
MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
|
||||
ASSERT_EQ(ret->size(), 1);
|
||||
|
||||
ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
|
||||
const auto& [inner_data0, children0] =
|
||||
ret->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
|
||||
ASSERT_TRUE((children0->size() == 2));
|
||||
|
||||
ASSERT_TRUE((children0->at(0).Has<IntTreeInnerT>()));
|
||||
const auto& [inner_data_left0, children_left0] =
|
||||
children0->at(0).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data_left0 == IntTreeLeafT{4, 5}));
|
||||
ASSERT_TRUE((children_left0->size() == 1));
|
||||
ASSERT_TRUE((children_left0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children_left0->at(0).Get<IntTreeLeafT>() ==
|
||||
IntTreeLeafT{1, 2, 3, 4, 5}));
|
||||
|
||||
ASSERT_TRUE((children0->at(1).Has<IntTreeInnerT>()));
|
||||
const auto& [inner_data_right0, children_right0] =
|
||||
children0->at(1).Get<IntTreeInnerT>().tuple();
|
||||
ASSERT_TRUE((inner_data_right0 == IntTreeLeafT{6, 7}));
|
||||
ASSERT_TRUE((children_right0->size() == 1));
|
||||
ASSERT_TRUE((children_right0->at(0).Has<IntTreeLeafT>()));
|
||||
ASSERT_TRUE((children_right0->at(0).Get<IntTreeLeafT>() ==
|
||||
IntTreeLeafT{1, 2, 3, 6, 7}));
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace cinn::adt
|
||||
@@ -0,0 +1 @@
|
||||
cinn_cc_test(test_tensor_group SRCS tensor_group_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,138 @@
|
||||
// Copyright (c) 2023 CINN 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 <gtest/gtest.h>
|
||||
#include <vector>
|
||||
#include "paddle/utils/flat_hash_map.h"
|
||||
|
||||
#include "paddle/cinn/ast_gen_ius/tensor_group.h"
|
||||
#include "paddle/cinn/ir/ir.h"
|
||||
#include "paddle/cinn/ir/ir_base.h"
|
||||
#include "paddle/cinn/ir/tensor.h"
|
||||
#include "paddle/cinn/lang/compute.h"
|
||||
#include "paddle/cinn/lang/placeholder.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace ast_gen_ius {
|
||||
|
||||
using ir::Expr;
|
||||
using ir::Tensor;
|
||||
using ir::Var;
|
||||
using lang::Compute;
|
||||
using lang::Placeholder;
|
||||
|
||||
TEST(TensorGroup, Easy) {
|
||||
auto M = Expr(100);
|
||||
auto N = Expr(15);
|
||||
Placeholder<float> A("A", {M, N});
|
||||
|
||||
Tensor B = Compute(
|
||||
{M, N}, [=](Var i, Var j) -> Expr { return A(i, j) + 1.f; }, "B");
|
||||
|
||||
TensorGroup tensor_group({B});
|
||||
|
||||
ASSERT_TRUE(tensor_group.Contain("A"));
|
||||
ASSERT_TRUE(tensor_group.Contain("B"));
|
||||
ASSERT_EQ(tensor_group.Get("B")->name, "B");
|
||||
ASSERT_EQ(tensor_group.Get("A")->name, "A");
|
||||
ASSERT_EQ(tensor_group.GetAllTensors().size(), 2UL);
|
||||
|
||||
ASSERT_EQ(tensor_group.GetCtrlDepTensors("A").size(), 0UL);
|
||||
ASSERT_EQ(tensor_group.GetCtrlDepTensors("B").size(), 1UL);
|
||||
ASSERT_TRUE(tensor_group.GetCtrlDepTensors("B").count(A));
|
||||
|
||||
std::vector<ir::Tensor> topo_tensors =
|
||||
tensor_group.GetGenFuncTopoOrder({A.tensor(), B});
|
||||
ASSERT_EQ(topo_tensors.size(), 1UL);
|
||||
ASSERT_EQ(topo_tensors[0]->name, "B");
|
||||
|
||||
ASSERT_EQ(tensor_group.GetShareMemRootName("A"), "A");
|
||||
ASSERT_EQ(tensor_group.GetShareMemRootName("B"), "B");
|
||||
tensor_group.MarkShareMemBuffer(tensor_group.Get("A"), tensor_group.Get("B"));
|
||||
|
||||
paddle::flat_hash_map<std::string, ir::Tensor> buffered_tensors =
|
||||
tensor_group.AllocateBuffers();
|
||||
ASSERT_EQ(buffered_tensors["A"]->buffer->name,
|
||||
buffered_tensors["B"]->buffer->name);
|
||||
}
|
||||
|
||||
TEST(TensorGroup, GraphTopo) {
|
||||
auto M = Expr(16);
|
||||
auto N = Expr(16);
|
||||
|
||||
/*
|
||||
* A B
|
||||
* / \ /
|
||||
* C D
|
||||
* \ /
|
||||
* E
|
||||
*/
|
||||
|
||||
Placeholder<float> A("A", {M, N});
|
||||
Placeholder<float> B("B", {M, N});
|
||||
|
||||
Tensor C = Compute(
|
||||
{M, N}, [=](Var i, Var j) -> Expr { return A(i, j) + 1.f; }, "C");
|
||||
|
||||
Tensor D = Compute(
|
||||
{M, N}, [=](Var i, Var j) -> Expr { return A(i, j) + B(i, j); }, "D");
|
||||
|
||||
Tensor E = Compute(
|
||||
{M, N}, [=](Var i, Var j) -> Expr { return C(i, j) / D(i, j); }, "E");
|
||||
|
||||
TensorGroup tensor_group({C, D, E});
|
||||
|
||||
std::vector<std::string> check_names = {"A", "B", "C", "D", "E"};
|
||||
ASSERT_EQ(tensor_group.GetAllTensors().size(), check_names.size());
|
||||
for (const std::string& name : check_names) {
|
||||
ASSERT_TRUE(tensor_group.Contain(name));
|
||||
ASSERT_EQ(tensor_group.Get(name)->name, name);
|
||||
}
|
||||
|
||||
ASSERT_TRUE(tensor_group.GetCtrlDepTensors("E").count(D));
|
||||
ASSERT_TRUE(tensor_group.GetCtrlDepTensors("E").count(C));
|
||||
ASSERT_TRUE(tensor_group.GetCtrlDepTensors("D").count(A));
|
||||
ASSERT_TRUE(tensor_group.GetCtrlDepTensors("D").count(B));
|
||||
ASSERT_TRUE(tensor_group.GetCtrlDepTensors("C").count(A));
|
||||
|
||||
std::vector<ir::Tensor> topo_tensors = tensor_group.GetGenFuncTopoOrder();
|
||||
ASSERT_EQ(topo_tensors.size(), check_names.size());
|
||||
for (size_t i = 0; i < check_names.size(); ++i) {
|
||||
ASSERT_EQ(topo_tensors[i]->name, check_names[i]);
|
||||
}
|
||||
|
||||
std::vector<ir::Tensor> topo_except_argu =
|
||||
tensor_group.GetGenFuncTopoOrder({A.tensor(), B.tensor()});
|
||||
ASSERT_EQ(topo_except_argu.size(), 3);
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
ASSERT_EQ(topo_except_argu[i]->name, check_names[i + 2]);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < check_names.size(); ++i) {
|
||||
ASSERT_EQ(tensor_group.GetShareMemRootName(check_names[i]), check_names[i]);
|
||||
}
|
||||
tensor_group.MarkShareMemBuffer(tensor_group.Get("A"), tensor_group.Get("B"));
|
||||
tensor_group.MarkShareMemBuffer(tensor_group.Get("B"), tensor_group.Get("C"));
|
||||
tensor_group.MarkShareMemBuffer(tensor_group.Get("C"), tensor_group.Get("D"));
|
||||
|
||||
ASSERT_EQ(tensor_group.GetShareMemRootName("A"),
|
||||
tensor_group.GetShareMemRootName("D"));
|
||||
paddle::flat_hash_map<std::string, ir::Tensor> buffered_tensors =
|
||||
tensor_group.AllocateBuffers();
|
||||
ASSERT_EQ(buffered_tensors["A"]->buffer->name,
|
||||
buffered_tensors["D"]->buffer->name);
|
||||
}
|
||||
|
||||
} // namespace ast_gen_ius
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,23 @@
|
||||
cinn_cc_test(test_dfs_walker SRCS dfs_walker_test.cc DEPS gtest glog)
|
||||
cinn_cc_test(test_dfs_topo_walker SRCS dfs_topo_walker_test.cc DEPS gtest glog)
|
||||
cinn_cc_test(test_cinn_value SRCS cinn_value_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_axis SRCS axis_test.cc DEPS cinncore)
|
||||
|
||||
cinn_cc_test(dim_expr_converter_test SRCS dim_expr_converter_test.cc DEPS
|
||||
cinncore)
|
||||
cinn_cc_test(broadcast_tree_test SRCS broadcast_tree_test.cc DEPS cinncore)
|
||||
|
||||
cinn_cc_test(test_equation_graph_topo_walker SRCS
|
||||
equation_graph_topo_walker_test.cc DEPS gtest glog)
|
||||
cinn_cc_test(test_type SRCS type_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_topo_walker SRCS topo_walker_test.cc DEPS gtest glog)
|
||||
cinn_cc_test(test_shared SRCS shared_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_is_reachable_predicator SRCS is_reachable_predicator_test.cc
|
||||
DEPS gtest glog)
|
||||
cinn_cc_test(test_integer_set SRCS integer_set_test.cc DEPS cinncore)
|
||||
if(WITH_CUDA)
|
||||
cinn_nv_test(test_fp16_bf16_cuda SRCS float16_bfloat16_cuda_test.cu DEPS
|
||||
gtest glog)
|
||||
endif()
|
||||
cinn_cc_test(test_fp16_bf16_host SRCS float16_bfloat16_host_test.cc DEPS gtest
|
||||
glog)
|
||||
@@ -0,0 +1,45 @@
|
||||
// Copyright (c) 2023 CINN 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 "paddle/cinn/common/axis.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "paddle/cinn/utils/string.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
TEST(AXISNAME, BASE) {
|
||||
ASSERT_EQ(axis_name(0), std::string("i"));
|
||||
ASSERT_EQ(axis_name(1), std::string("j"));
|
||||
ASSERT_EQ(axis_name(22), std::string("ii"));
|
||||
ASSERT_EQ(axis_name(44), std::string("iii"));
|
||||
}
|
||||
|
||||
TEST(AXISNAME, CHECK_RESERVED) {
|
||||
ASSERT_TRUE(IsAxisNameReserved("i"));
|
||||
ASSERT_TRUE(IsAxisNameReserved("j"));
|
||||
ASSERT_TRUE(IsAxisNameReserved("ii"));
|
||||
ASSERT_TRUE(IsAxisNameReserved("iiiiiiiiii"));
|
||||
ASSERT_FALSE(IsAxisNameReserved("ijk"));
|
||||
ASSERT_FALSE(IsAxisNameReserved("iiiiiiiiiij"));
|
||||
ASSERT_FALSE(IsAxisNameReserved("x"));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,118 @@
|
||||
// Copyright (c) 2024 PaddlePaddle 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 "paddle/cinn/common/broadcast_tree.h"
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace cinn::common {
|
||||
using namespace symbol; // NOLINT
|
||||
|
||||
namespace {
|
||||
|
||||
DimExpr MakeBroadcastDimExpr(const DimExpr& expr1, const DimExpr& expr2) {
|
||||
List<DimExpr> operands{expr1, expr2};
|
||||
return Broadcast<DimExpr>{operands};
|
||||
}
|
||||
|
||||
bool DimExprNonBroadcast(const DimExpr& dim_expr) {
|
||||
if (dim_expr.Has<Broadcast<DimExpr>>()) {
|
||||
return false;
|
||||
} else {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
void CheckLeafNonBroadcast(const BroadcastLeaf& leaf) {
|
||||
for (const auto& operands : *leaf) {
|
||||
for (const auto& operand : operands) {
|
||||
ASSERT_TRUE(DimExprNonBroadcast(operand));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void CheckInnerBranchNonBroadcast(
|
||||
const BroadcastBranch<BroadcastTree>& branch) {
|
||||
const auto& [_, lhs_eq_rhs_tree, lhs_eq_one_tree, rhs_eq_one_tree] =
|
||||
branch.tuple();
|
||||
ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastLeaf>());
|
||||
ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastLeaf>());
|
||||
ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastLeaf>());
|
||||
CheckLeafNonBroadcast(lhs_eq_rhs_tree.Get<BroadcastLeaf>());
|
||||
CheckLeafNonBroadcast(lhs_eq_one_tree.Get<BroadcastLeaf>());
|
||||
CheckLeafNonBroadcast(rhs_eq_one_tree.Get<BroadcastLeaf>());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
TEST(BroadcastTree, Naive) {
|
||||
DimExpr expr1("S1");
|
||||
DimExpr expr2("S2");
|
||||
DimExpr expr3("S3");
|
||||
DimExpr expr4("S4");
|
||||
std::vector<DimExpr> tensor_shape{expr1,
|
||||
expr2,
|
||||
MakeBroadcastDimExpr(expr1, expr2),
|
||||
MakeBroadcastDimExpr(expr3, expr4)};
|
||||
BroadcastLeaf leaf = adt::List<std::vector<DimExpr>>{tensor_shape};
|
||||
int num_of_leaves = 0;
|
||||
BroadcastTree tree = ConstructBroadcastTree(leaf, &num_of_leaves);
|
||||
ASSERT_TRUE(tree.Has<BroadcastBranch<BroadcastTree>>());
|
||||
const auto& branch = tree.Get<BroadcastBranch<BroadcastTree>>();
|
||||
const auto& [cstr_broadcastable,
|
||||
lhs_eq_rhs_tree,
|
||||
lhs_eq_one_tree,
|
||||
rhs_eq_one_tree] = branch.tuple();
|
||||
ASSERT_EQ(cstr_broadcastable->lhs, DimExpr("S1"));
|
||||
ASSERT_EQ(cstr_broadcastable->rhs, DimExpr("S2"));
|
||||
ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastBranch<BroadcastTree>>());
|
||||
ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastBranch<BroadcastTree>>());
|
||||
ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastBranch<BroadcastTree>>());
|
||||
CheckInnerBranchNonBroadcast(
|
||||
lhs_eq_rhs_tree.Get<BroadcastBranch<BroadcastTree>>());
|
||||
CheckInnerBranchNonBroadcast(
|
||||
lhs_eq_one_tree.Get<BroadcastBranch<BroadcastTree>>());
|
||||
CheckInnerBranchNonBroadcast(
|
||||
rhs_eq_one_tree.Get<BroadcastBranch<BroadcastTree>>());
|
||||
}
|
||||
|
||||
TEST(BroadcastTree, SimplifyConstantBroadcast) {
|
||||
DimExpr expr1("S1");
|
||||
DimExpr expr2("S2");
|
||||
DimExpr expr3("S3");
|
||||
DimExpr expr4(4);
|
||||
std::vector<DimExpr> tensor_shape{expr1,
|
||||
expr2,
|
||||
MakeBroadcastDimExpr(expr1, expr2),
|
||||
MakeBroadcastDimExpr(expr3, expr4)};
|
||||
BroadcastLeaf leaf = adt::List<std::vector<DimExpr>>{tensor_shape};
|
||||
int num_of_leaves = 0;
|
||||
BroadcastTree tree = ConstructBroadcastTree(leaf, &num_of_leaves);
|
||||
ASSERT_TRUE(tree.Has<BroadcastBranch<BroadcastTree>>());
|
||||
const auto& branch = tree.Get<BroadcastBranch<BroadcastTree>>();
|
||||
const auto& [cstr_broadcastable,
|
||||
lhs_eq_rhs_tree,
|
||||
lhs_eq_one_tree,
|
||||
rhs_eq_one_tree] = branch.tuple();
|
||||
ASSERT_EQ(cstr_broadcastable->lhs, DimExpr("S1"));
|
||||
ASSERT_EQ(cstr_broadcastable->rhs, DimExpr("S2"));
|
||||
ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastLeaf>());
|
||||
ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastLeaf>());
|
||||
ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastLeaf>());
|
||||
CheckLeafNonBroadcast(lhs_eq_rhs_tree.Get<BroadcastLeaf>());
|
||||
CheckLeafNonBroadcast(lhs_eq_one_tree.Get<BroadcastLeaf>());
|
||||
CheckLeafNonBroadcast(rhs_eq_one_tree.Get<BroadcastLeaf>());
|
||||
}
|
||||
|
||||
} // namespace cinn::common
|
||||
@@ -0,0 +1,59 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/common/cinn_value.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/common/common.h"
|
||||
#include "paddle/cinn/common/ir_util.h"
|
||||
#include "paddle/cinn/ir/ir.h"
|
||||
#include "paddle/cinn/ir/ir_printer.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
TEST(CINNValue, test) {
|
||||
{
|
||||
CINNValue value(32);
|
||||
ASSERT_EQ(int(value), 32); // NOLINT
|
||||
}
|
||||
{
|
||||
CINNValue value(32.f);
|
||||
ASSERT_NEAR(float(value), 32.f, 1e-6); // NOLINT
|
||||
}
|
||||
}
|
||||
|
||||
TEST(CINNValue, buffer) {
|
||||
cinn_buffer_t* v = nullptr;
|
||||
CINNValue value(v);
|
||||
ASSERT_EQ((cinn_buffer_t*)value, nullptr);
|
||||
}
|
||||
|
||||
TEST(CINNValue, Expr) {
|
||||
Expr a(1);
|
||||
|
||||
{
|
||||
CINNValue value(a);
|
||||
ASSERT_TRUE(a == value);
|
||||
}
|
||||
|
||||
{
|
||||
CINNValue copied = CINNValue(a);
|
||||
ASSERT_TRUE(copied == cinn::common::make_const(1));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,54 @@
|
||||
// Copyright (c) 2023 CINN 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 <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/common/dfs_topo_walker.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
TEST(DfsTopoWalker, simple) {
|
||||
std::vector<std::pair<int, int>> edges{
|
||||
{0, 1}, {2, 3}, {1, 3}, {0, 3}, {3, 4}};
|
||||
DfsTopoWalker<int> walker(
|
||||
[&](int node, const std::function<void(int)>& NodeHandler) {
|
||||
for (const auto& pair : edges) {
|
||||
if (pair.second == node) {
|
||||
NodeHandler(pair.first);
|
||||
}
|
||||
}
|
||||
},
|
||||
[&](int node, const std::function<void(int)>& NodeHandler) {
|
||||
for (const auto& pair : edges) {
|
||||
if (pair.first == node) {
|
||||
NodeHandler(pair.second);
|
||||
}
|
||||
}
|
||||
});
|
||||
std::vector<int> sources{0, 2};
|
||||
std::vector<int> outputs;
|
||||
walker(sources.begin(), sources.end(), [&](int node) {
|
||||
outputs.push_back(node);
|
||||
});
|
||||
for (auto output : outputs) {
|
||||
LOG(INFO) << output;
|
||||
}
|
||||
std::vector<int> expected{0, 1, 2, 3, 4};
|
||||
EXPECT_TRUE((outputs == expected));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,72 @@
|
||||
// Copyright (c) 2023 CINN 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 "paddle/cinn/common/dfs_walker.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
TEST(DfsWalker, simple_on_push) {
|
||||
DfsWalker<int> visitor(
|
||||
[](int node, const std::function<void(int)>& NodeHandler) {
|
||||
if (node == 0) {
|
||||
NodeHandler(3);
|
||||
} else if (node == 1) {
|
||||
NodeHandler(2);
|
||||
NodeHandler(3);
|
||||
} else if (node == 2 || node == 3) {
|
||||
NodeHandler(4);
|
||||
}
|
||||
});
|
||||
std::vector<int> sources{0, 1};
|
||||
std::vector<int> outputs;
|
||||
visitor(sources.begin(), sources.end(), [&](int node) {
|
||||
LOG(ERROR) << node;
|
||||
outputs.push_back(node);
|
||||
});
|
||||
std::vector<int> expected{0, 3, 4, 1, 2};
|
||||
EXPECT_TRUE((outputs == expected));
|
||||
}
|
||||
|
||||
TEST(DfsWalker, simple_on_pop) {
|
||||
DfsWalker<int> visitor(
|
||||
[](int node, const std::function<void(int)>& NodeHandler) {
|
||||
if (node == 0) {
|
||||
NodeHandler(3);
|
||||
} else if (node == 1) {
|
||||
NodeHandler(2);
|
||||
NodeHandler(3);
|
||||
} else if (node == 2 || node == 3) {
|
||||
NodeHandler(4);
|
||||
}
|
||||
});
|
||||
std::vector<int> sources{0, 1};
|
||||
std::vector<int> outputs;
|
||||
visitor(
|
||||
sources.begin(),
|
||||
sources.end(),
|
||||
[](int) {},
|
||||
[&](int node) {
|
||||
LOG(ERROR) << node;
|
||||
outputs.push_back(node);
|
||||
});
|
||||
std::vector<int> expected{4, 3, 0, 2, 1};
|
||||
EXPECT_TRUE((outputs == expected));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,96 @@
|
||||
// Copyright (c) 2023 PaddlePaddle 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 <sstream>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#include "paddle/cinn/common/dim_expr_converter.h"
|
||||
#include "paddle/cinn/common/ir_util.h"
|
||||
#include "paddle/cinn/ir/ir_printer.h"
|
||||
|
||||
namespace cinn::common::test {
|
||||
|
||||
using namespace symbol; // NOLINT
|
||||
|
||||
TEST(Convert, AddExpr) {
|
||||
List<DimExpr> num_lists{DimExpr(4), DimExpr(5), DimExpr("sym_0")};
|
||||
DimExpr dim_expr{Add<DimExpr>{num_lists}};
|
||||
ir::Expr src_expr = DimExprConverter().ConvertToIrExpr(dim_expr);
|
||||
|
||||
ir::Expr expr1 =
|
||||
ir::Add::Make(ir::Expr(std::int64_t(4)), ir::Expr(std::int64_t(5)));
|
||||
ir::Expr dst_expr =
|
||||
ir::Add::Make(expr1,
|
||||
ir::_Var_::Make(ir::Expr(static_cast<int64_t>(1)),
|
||||
ir::Expr(INT32_MAX),
|
||||
"sym_0",
|
||||
/* is_reduce = */ false,
|
||||
/* is_symbolic_constant = */ true));
|
||||
ASSERT_TRUE(MathEqual(src_expr, dst_expr));
|
||||
}
|
||||
|
||||
TEST(Convert, SubExpr) {
|
||||
DimExpr dim_expr = DimExpr(4) - DimExpr("sym_0");
|
||||
ir::Expr src_expr = DimExprConverter().ConvertToIrExpr(dim_expr);
|
||||
|
||||
ir::Expr expr1 =
|
||||
ir::Sub::Make(ir::Expr(std::int64_t(0)),
|
||||
ir::_Var_::Make(ir::Expr(static_cast<int64_t>(1)),
|
||||
ir::Expr(INT32_MAX),
|
||||
"sym_0",
|
||||
/* is_reduce = */ false,
|
||||
/* is_symbolic_constant = */ true));
|
||||
ir::Expr dst_expr = ir::Add::Make(ir::Expr(std::int64_t(4)), expr1);
|
||||
ASSERT_TRUE(MathEqual(src_expr, dst_expr));
|
||||
}
|
||||
|
||||
TEST(Convert, MulExpr) {
|
||||
List<DimExpr> num_lists{DimExpr(4), DimExpr(5), DimExpr("sym_0")};
|
||||
DimExpr dim_expr{Mul<DimExpr>{num_lists}};
|
||||
ir::Expr src_expr = DimExprConverter().ConvertToIrExpr(dim_expr);
|
||||
|
||||
ir::Expr expr1 =
|
||||
ir::Mul::Make(ir::Expr(std::int64_t(4)), ir::Expr(std::int64_t(5)));
|
||||
ir::Expr dst_expr =
|
||||
ir::Mul::Make(expr1,
|
||||
ir::_Var_::Make(ir::Expr(static_cast<int64_t>(1)),
|
||||
ir::Expr(INT32_MAX),
|
||||
"sym_0",
|
||||
/* is_reduce = */ false,
|
||||
/* is_symbolic_constant = */ true));
|
||||
ASSERT_TRUE(MathEqual(src_expr, dst_expr));
|
||||
}
|
||||
|
||||
TEST(Convert, MaxExpr) {
|
||||
List<DimExpr> num_lists{DimExpr(4), DimExpr(5), DimExpr("sym_0")};
|
||||
DimExpr dim_expr{Max<DimExpr>{num_lists}};
|
||||
ir::Expr src_expr = DimExprConverter().ConvertToIrExpr(dim_expr);
|
||||
|
||||
std::ostringstream stream;
|
||||
stream << src_expr;
|
||||
ASSERT_EQ(stream.str(), "cinn_max(cinn_max(4ll, 5ll), sym_0)");
|
||||
}
|
||||
|
||||
TEST(Convert, MinExpr) {
|
||||
List<DimExpr> num_lists{DimExpr(4), DimExpr(5), DimExpr("sym_0")};
|
||||
DimExpr dim_expr{Min<DimExpr>{num_lists}};
|
||||
ir::Expr src_expr = DimExprConverter().ConvertToIrExpr(dim_expr);
|
||||
|
||||
std::ostringstream stream;
|
||||
stream << src_expr;
|
||||
ASSERT_EQ(stream.str(), "cinn_min(cinn_min(4ll, 5ll), sym_0)");
|
||||
}
|
||||
|
||||
} // namespace cinn::common::test
|
||||
@@ -0,0 +1,116 @@
|
||||
// Copyright (c) 2023 PaddlePaddle 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.
|
||||
|
||||
// TODO(yifan): Add unittest here
|
||||
#include "paddle/cinn/common/equation_graph_topo_walker.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace adt {
|
||||
namespace common {
|
||||
|
||||
using VT = int;
|
||||
using FT = std::string;
|
||||
/*
|
||||
Graph ex:
|
||||
|
||||
1-> "1->10" -> 10
|
||||
2-> "2->20" -> 20
|
||||
*/
|
||||
|
||||
TEST(EquationGraphTopoWalker, simple1) {
|
||||
auto F4V = [](VT variable, const std::function<void(FT)>& visitor) {
|
||||
if (variable == 1) {
|
||||
visitor("1->10");
|
||||
} else if (variable == 2) {
|
||||
visitor("2->20");
|
||||
}
|
||||
};
|
||||
auto InV4F = [](FT function, const std::function<void(VT)>& visitor) {
|
||||
if (function == "1->10") {
|
||||
visitor(1);
|
||||
} else if (function == "2->20") {
|
||||
visitor(2);
|
||||
}
|
||||
};
|
||||
auto OutV4F = [](FT function, const std::function<void(VT)>& visitor) {
|
||||
if (function == "1->10") {
|
||||
visitor(10);
|
||||
} else if (function == "2->20") {
|
||||
visitor(20);
|
||||
}
|
||||
};
|
||||
cinn::EquationGraphTopoWalker<VT, FT> walker(F4V, InV4F, OutV4F);
|
||||
std::vector<FT> outputs;
|
||||
std::function<void(FT)> FunctionVisitor = [&](FT function) {
|
||||
outputs.push_back(function);
|
||||
};
|
||||
walker.WalkFunction(1, FunctionVisitor);
|
||||
|
||||
std::vector<FT> expected{"1->10"};
|
||||
EXPECT_TRUE((outputs == expected));
|
||||
}
|
||||
|
||||
/*
|
||||
Graph ex:
|
||||
|
||||
1 -> "1->10, 1->11" -> 10
|
||||
-> 11
|
||||
2 -> "2->20" -> 20
|
||||
3 -> "3->30, 3->31" -> 30
|
||||
-> 31
|
||||
*/
|
||||
TEST(EquationGraphTopoWalker, simple2) {
|
||||
auto F4V = [](VT variable, const std::function<void(FT)>& visitor) {
|
||||
if (variable == 1) {
|
||||
visitor("1->10, 1->11");
|
||||
} else if (variable == 2) {
|
||||
visitor("2->20");
|
||||
} else if (variable == 3) {
|
||||
visitor("3->30, 3->31");
|
||||
}
|
||||
};
|
||||
auto InV4F = [](FT function, const std::function<void(VT)>& visitor) {
|
||||
if (function == "1->10, 1->11") {
|
||||
visitor(1);
|
||||
} else if (function == "2->20") {
|
||||
visitor(2);
|
||||
} else if (function == "3->30, 3->31") {
|
||||
visitor(3);
|
||||
}
|
||||
};
|
||||
auto OutV4F = [](FT function, const std::function<void(VT)>& visitor) {
|
||||
if (function == "1->10, 1->11") {
|
||||
visitor(10);
|
||||
visitor(11);
|
||||
} else if (function == "2->20") {
|
||||
visitor(20);
|
||||
} else if (function == "3->30, 3->31") {
|
||||
visitor(30);
|
||||
visitor(31);
|
||||
}
|
||||
};
|
||||
cinn::EquationGraphTopoWalker<VT, FT> walker(F4V, InV4F, OutV4F);
|
||||
std::vector<VT> outputs;
|
||||
std::function<void(VT)> VariableVisitor = [&](VT variable) {
|
||||
outputs.push_back(variable);
|
||||
};
|
||||
walker.WalkVariable(1, VariableVisitor);
|
||||
std::vector<VT> expected{1, 10, 11};
|
||||
EXPECT_TRUE((outputs == expected));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace adt
|
||||
@@ -0,0 +1,286 @@
|
||||
// Copyright (c) 2021 CINN 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 <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <random>
|
||||
#include <vector>
|
||||
#include "paddle/cinn/common/bfloat16.h"
|
||||
#include "paddle/cinn/common/float16.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
#define CUDA_CALL(func) \
|
||||
{ \
|
||||
auto status = func; \
|
||||
if (status != cudaSuccess) { \
|
||||
std::stringstream ss; \
|
||||
ss << "CUDA Error : " << cudaGetErrorString(status); \
|
||||
PADDLE_THROW(::common::errors::Fatal(ss.str())); \
|
||||
} \
|
||||
}
|
||||
|
||||
class CudaMem {
|
||||
public:
|
||||
CudaMem() = default;
|
||||
|
||||
void* mutable_data(size_t bytes) {
|
||||
PADDLE_ENFORCE_GT(
|
||||
bytes,
|
||||
0,
|
||||
::common::errors::InvalidArgument("Cannot allocate empty memory!"));
|
||||
if (ptr) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
bytes,
|
||||
bytes_,
|
||||
::common::errors::InvalidArgument("Try allocate memory twice!"));
|
||||
return ptr;
|
||||
}
|
||||
CUDA_CALL(cudaMalloc(&ptr, bytes));
|
||||
bytes_ = bytes;
|
||||
return ptr;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* mutable_data(size_t num) {
|
||||
return reinterpret_cast<T*>(mutable_data(num * sizeof(T)));
|
||||
}
|
||||
|
||||
void* data() const {
|
||||
PADDLE_ENFORCE_NOT_NULL(ptr,
|
||||
::common::errors::InvalidArgument(
|
||||
"Pointer is null; please ensure it is properly "
|
||||
"initialized before use."));
|
||||
return ptr;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* data() const {
|
||||
return reinterpret_cast<T*>(data());
|
||||
}
|
||||
|
||||
void MemcpyFromHost(const void* src,
|
||||
size_t bytes,
|
||||
cudaStream_t stream = nullptr) {
|
||||
PADDLE_ENFORCE_LE(
|
||||
bytes,
|
||||
bytes_,
|
||||
::common::errors::InvalidArgument("Too many data need copy"));
|
||||
CUDA_CALL(cudaMemcpyAsync(ptr, src, bytes, cudaMemcpyHostToDevice, stream));
|
||||
}
|
||||
|
||||
void MemcpyToHost(void* dst, size_t bytes, cudaStream_t stream = nullptr) {
|
||||
PADDLE_ENFORCE_LE(
|
||||
bytes,
|
||||
bytes_,
|
||||
::common::errors::InvalidArgument("Too many data need copy"));
|
||||
CUDA_CALL(cudaMemcpyAsync(dst, ptr, bytes, cudaMemcpyDeviceToHost, stream));
|
||||
}
|
||||
|
||||
~CudaMem() {
|
||||
if (ptr) {
|
||||
cudaFree(ptr);
|
||||
}
|
||||
bytes_ = 0;
|
||||
}
|
||||
|
||||
private:
|
||||
void* ptr{nullptr};
|
||||
size_t bytes_{0};
|
||||
};
|
||||
|
||||
__global__ void cast_fp32_to_fp16_cuda_kernel(const float* input,
|
||||
const int num,
|
||||
float16* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
out[idx] = float16(input[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void cast_fp16_to_fp32_cuda_kernel(const float16* input,
|
||||
const int num,
|
||||
float* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
out[idx] = static_cast<float>(input[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void test_fp16_cuda_kernel(const float16* x,
|
||||
const float16* y,
|
||||
const int num,
|
||||
float16* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
float16 x_i = x[idx], y_i = y[idx];
|
||||
x_i += float16(1);
|
||||
|
||||
out[idx] = (x_i + y_i) * (x_i - y_i);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void cast_fp32_to_bf16_cuda_kernel(const float* input,
|
||||
const int num,
|
||||
bfloat16* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
out[idx] = bfloat16(input[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void cast_bf16_to_fp32_cuda_kernel(const bfloat16* input,
|
||||
const int num,
|
||||
float* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
out[idx] = static_cast<float>(input[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void test_bf16_cuda_kernel(const bfloat16* x,
|
||||
const bfloat16* y,
|
||||
const int num,
|
||||
bfloat16* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
bfloat16 x_i = x[idx], y_i = y[idx];
|
||||
x_i += bfloat16(1);
|
||||
|
||||
out[idx] = (x_i + y_i) * (x_i - y_i);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void test_fp32_cuda_kernel(const float* x,
|
||||
const float* y,
|
||||
const int num,
|
||||
float* out) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
float x_i = x[idx], y_i = y[idx];
|
||||
x_i += 1.0f;
|
||||
|
||||
out[idx] = (x_i + y_i) * (x_i - y_i);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(FP16_BF16, basic_cuda) {
|
||||
#ifdef CUDA_VERSION
|
||||
LOG(INFO) << "CUDA version: " << CUDA_VERSION;
|
||||
#endif
|
||||
|
||||
int num = 2048;
|
||||
|
||||
cudaStream_t stream;
|
||||
CUDA_CALL(cudaStreamCreate(&stream));
|
||||
|
||||
dim3 block = 1024;
|
||||
dim3 grid = (num + block.x - 1) / block.x;
|
||||
|
||||
std::vector<float> x_fp32_host(num), y_fp32_host(num);
|
||||
{ // step1 : generate input data
|
||||
std::random_device r;
|
||||
std::default_random_engine eng(r());
|
||||
std::uniform_real_distribution<float> dis(1e-5f, 1.0f);
|
||||
|
||||
for (int i = 0; i < num; ++i) {
|
||||
x_fp32_host[i] = dis(eng);
|
||||
y_fp32_host[i] = dis(eng);
|
||||
}
|
||||
}
|
||||
|
||||
CudaMem x_fp32_device, y_fp32_device, out_fp32_device;
|
||||
{ // step2 : compute fp32 result
|
||||
auto x_fp32_ptr = x_fp32_device.mutable_data<float>(num);
|
||||
auto y_fp32_ptr = y_fp32_device.mutable_data<float>(num);
|
||||
auto out_fp32_ptr = out_fp32_device.mutable_data<float>(num);
|
||||
|
||||
x_fp32_device.MemcpyFromHost(
|
||||
x_fp32_host.data(), num * sizeof(float), stream);
|
||||
y_fp32_device.MemcpyFromHost(
|
||||
y_fp32_host.data(), num * sizeof(float), stream);
|
||||
|
||||
test_fp32_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
x_fp32_ptr, y_fp32_ptr, num, out_fp32_ptr);
|
||||
}
|
||||
|
||||
CudaMem x_fp16_device, y_fp16_device, out_fp16_device;
|
||||
CudaMem x_bf16_device, y_bf16_device, out_bf16_device;
|
||||
{ // step3 : compute fp16/bf16 result
|
||||
// step3.1 : compute fp16 result
|
||||
auto x_fp16_ptr = x_fp16_device.mutable_data<float16>(num);
|
||||
auto y_fp16_ptr = y_fp16_device.mutable_data<float16>(num);
|
||||
auto out_fp16_ptr = out_fp16_device.mutable_data<float16>(num);
|
||||
|
||||
cast_fp32_to_fp16_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
x_fp32_device.data<float>(), num, x_fp16_ptr);
|
||||
cast_fp32_to_fp16_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
y_fp32_device.data<float>(), num, y_fp16_ptr);
|
||||
|
||||
test_fp16_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
x_fp16_ptr, y_fp16_ptr, num, out_fp16_ptr);
|
||||
|
||||
// step3.2 : compute bf16 result
|
||||
auto x_bf16_ptr = x_bf16_device.mutable_data<bfloat16>(num);
|
||||
auto y_bf16_ptr = y_bf16_device.mutable_data<bfloat16>(num);
|
||||
auto out_bf16_ptr = out_bf16_device.mutable_data<bfloat16>(num);
|
||||
|
||||
cast_fp32_to_bf16_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
x_fp32_device.data<float>(), num, x_bf16_ptr);
|
||||
cast_fp32_to_bf16_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
y_fp32_device.data<float>(), num, y_bf16_ptr);
|
||||
|
||||
test_bf16_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
x_bf16_ptr, y_bf16_ptr, num, out_bf16_ptr);
|
||||
}
|
||||
|
||||
CudaMem fp32res_fp16_device;
|
||||
CudaMem fp32res_bf16_device;
|
||||
{ // step4 : cast fp16/bf16 result to fp32 result
|
||||
// step4.1 : cast fp16 result to fp32 result
|
||||
auto fp32res_fp16_ptr = fp32res_fp16_device.mutable_data<float>(num);
|
||||
cast_fp16_to_fp32_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
out_fp16_device.data<float16>(), num, fp32res_fp16_ptr);
|
||||
|
||||
// step4.2 : cast bf16 result to fp32 result
|
||||
auto fp32res_bf16_ptr = fp32res_bf16_device.mutable_data<float>(num);
|
||||
cast_bf16_to_fp32_cuda_kernel<<<grid, block, 0, stream>>>(
|
||||
out_bf16_device.data<bfloat16>(), num, fp32res_bf16_ptr);
|
||||
}
|
||||
|
||||
std::vector<float> out_fp32_host(num), out_fp16_host(num), out_bf16_host(num);
|
||||
{ // step5 : copy result from device to host
|
||||
out_fp32_device.MemcpyToHost(
|
||||
out_fp32_host.data(), num * sizeof(float), stream);
|
||||
fp32res_fp16_device.MemcpyToHost(
|
||||
out_fp16_host.data(), num * sizeof(float), stream);
|
||||
fp32res_bf16_device.MemcpyToHost(
|
||||
out_bf16_host.data(), num * sizeof(float), stream);
|
||||
}
|
||||
|
||||
CUDA_CALL(cudaStreamSynchronize(stream));
|
||||
|
||||
for (int i = 0; i < num; ++i) {
|
||||
ASSERT_NEAR(out_fp32_host[i], out_fp16_host[i], 1e-2f);
|
||||
ASSERT_NEAR(out_fp32_host[i], out_bf16_host[i], 1e-1f);
|
||||
}
|
||||
|
||||
CUDA_CALL(cudaStreamDestroy(stream));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,104 @@
|
||||
// Copyright (c) 2021 CINN 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 <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/cinn/common/bfloat16.h"
|
||||
#include "paddle/cinn/common/float16.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
std::vector<float16> test_fp16_host_kernel(const float16* x,
|
||||
const float16* y,
|
||||
const int num) {
|
||||
std::vector<float16> out(num);
|
||||
for (int idx = 0; idx < num; ++idx) {
|
||||
float16 x_i = x[idx], y_i = y[idx];
|
||||
x_i += float16(1);
|
||||
|
||||
out[idx] = (x_i + y_i) * (x_i - y_i);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<bfloat16> test_bf16_host_kernel(const bfloat16* x,
|
||||
const bfloat16* y,
|
||||
const int num) {
|
||||
std::vector<bfloat16> out(num);
|
||||
for (int idx = 0; idx < num; ++idx) {
|
||||
bfloat16 x_i = x[idx], y_i = y[idx];
|
||||
x_i += bfloat16(1);
|
||||
|
||||
out[idx] = (x_i + y_i) * (x_i - y_i);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<float> test_fp32_host_kernel(const float* x,
|
||||
const float* y,
|
||||
const int num) {
|
||||
std::vector<float> out(num);
|
||||
for (int idx = 0; idx < num; ++idx) {
|
||||
float x_i = x[idx], y_i = y[idx];
|
||||
x_i += 1.0f;
|
||||
|
||||
out[idx] = (x_i + y_i) * (x_i - y_i);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
TEST(FP16_BF16, basic_host) {
|
||||
int num = 2048;
|
||||
// int num = 2;
|
||||
std::vector<float16> x_fp16(num), y_fp16(num);
|
||||
std::vector<bfloat16> x_bf16(num), y_bf16(num);
|
||||
std::vector<float> x_fp32(num), y_fp32(num);
|
||||
|
||||
std::random_device r;
|
||||
std::default_random_engine eng(r());
|
||||
std::uniform_real_distribution<float> dis(1e-5f, 1.0f);
|
||||
|
||||
for (int i = 0; i < num; ++i) {
|
||||
x_fp16[i] = x_fp32[i] = dis(eng);
|
||||
y_fp16[i] = y_fp32[i] = dis(eng);
|
||||
|
||||
x_fp16[i] = x_fp32[i];
|
||||
y_fp16[i] = y_fp32[i];
|
||||
|
||||
x_bf16[i] = x_fp32[i];
|
||||
y_bf16[i] = y_fp32[i];
|
||||
}
|
||||
|
||||
auto out_fp16 = test_fp16_host_kernel(x_fp16.data(), y_fp16.data(), num);
|
||||
ASSERT_EQ(out_fp16.size(), num);
|
||||
|
||||
auto out_bf16 = test_bf16_host_kernel(x_bf16.data(), y_bf16.data(), num);
|
||||
ASSERT_EQ(out_bf16.size(), num);
|
||||
|
||||
auto out_fp32 = test_fp32_host_kernel(x_fp32.data(), y_fp32.data(), num);
|
||||
ASSERT_EQ(out_fp32.size(), num);
|
||||
|
||||
for (int i = 0; i < num; ++i) {
|
||||
ASSERT_NEAR(static_cast<float>(out_fp16[i]), out_fp32[i], 1e-2f);
|
||||
ASSERT_NEAR(static_cast<float>(out_bf16[i]), out_fp32[i], 1e-1f);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,341 @@
|
||||
// Copyright (c) 2023 CINN 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 "paddle/cinn/common/integer_set.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
class TestSymbolicExprAnalyzer : public ::testing::Test {
|
||||
public:
|
||||
void SetUp() override {
|
||||
// Var is [lower_bound, upper_bound)
|
||||
i = ir::Var(ir::Expr(0), ir::Expr(7), "i"); // i ∈ [0, 7)
|
||||
j = ir::Var(ir::Expr(0), ir::Expr(15), "j"); // j ∈ [0, 15)
|
||||
// CasInterval is [lower_bound, upper_bound]
|
||||
var_intervals = {
|
||||
{"i", CasInterval(i->lower_bound, i->upper_bound - 1)}, // i ∈ [0, 6]
|
||||
{"j", CasInterval(j->lower_bound, j->upper_bound - 1)}, // j ∈ [0, 14]
|
||||
};
|
||||
}
|
||||
|
||||
ir::Var i;
|
||||
ir::Var j;
|
||||
cas_intervals_t var_intervals;
|
||||
SymbolicExprAnalyzer analyzer{var_intervals};
|
||||
};
|
||||
|
||||
TEST_F(TestSymbolicExprAnalyzer, bound) {
|
||||
ir::Expr e1 = i + j;
|
||||
EXPECT_EQ(analyzer.LowerBound(e1), ir::Expr(0));
|
||||
EXPECT_EQ(analyzer.UpperBound(e1), ir::Expr(20)); // 6 + 14 = 20
|
||||
|
||||
ir::Expr e2 = 16 * i + j;
|
||||
EXPECT_EQ(analyzer.LowerBound(e2), ir::Expr(0));
|
||||
EXPECT_EQ(analyzer.UpperBound(e2), ir::Expr(110)); // 16 * 6 + 14 = 110
|
||||
|
||||
ir::Expr e3 = 16 * i + j + 1;
|
||||
EXPECT_EQ(analyzer.LowerBound(e3), ir::Expr(1));
|
||||
EXPECT_EQ(analyzer.UpperBound(e3), ir::Expr(111)); // 16 * 6 + 15 = 111
|
||||
|
||||
ir::Expr e4 = (16 * i + j) / 16;
|
||||
EXPECT_EQ(analyzer.LowerBound(e4), ir::Expr(0));
|
||||
EXPECT_EQ(analyzer.UpperBound(e4), ir::Expr(6)); // 110 / 16 = 6
|
||||
|
||||
ir::Expr e5 = (16 * i + j) % 16;
|
||||
EXPECT_EQ(analyzer.LowerBound(e5), ir::Expr(0));
|
||||
EXPECT_EQ(analyzer.UpperBound(e5), ir::Expr(14)); // 110 % 16
|
||||
|
||||
ir::Expr e6 = i - j;
|
||||
EXPECT_EQ(analyzer.LowerBound(e6), ir::Expr(-14)); // 0 - 14
|
||||
EXPECT_EQ(analyzer.UpperBound(e6), ir::Expr(6)); // 6 - 0
|
||||
|
||||
ir::Expr e7 = 0 - i - j;
|
||||
EXPECT_EQ(analyzer.LowerBound(e7), ir::Expr(-20)); // 0 - 6 - 14
|
||||
EXPECT_EQ(analyzer.UpperBound(e7), ir::Expr(0)); // 0 - 0 - 0
|
||||
|
||||
ir::Expr e8 = -1 * i - j;
|
||||
EXPECT_EQ(analyzer.LowerBound(e8), ir::Expr(-20)); // -1 * 6 - 14
|
||||
EXPECT_EQ(analyzer.UpperBound(e8), ir::Expr(0)); // -1 * 0 - 0
|
||||
}
|
||||
|
||||
TEST_F(TestSymbolicExprAnalyzer, compare) {
|
||||
// case 1
|
||||
ir::Expr e1 = 4 * i + 2 * j;
|
||||
ir::Expr e2 = 2 * i + j;
|
||||
|
||||
EXPECT_TRUE(analyzer.ProveEQ(e1, e1).value() &&
|
||||
analyzer.Prove(ir::EQ::Make(e1, e1)).value());
|
||||
EXPECT_FALSE(analyzer.ProveEQ(e1, e2).has_value() ||
|
||||
analyzer.Prove(ir::EQ::Make(e1, e2)).has_value());
|
||||
EXPECT_FALSE(analyzer.ProveNE(e1, e1).value() &&
|
||||
analyzer.Prove(ir::NE::Make(e1, e1)).value());
|
||||
EXPECT_FALSE(analyzer.ProveNE(e1, e2).has_value() ||
|
||||
analyzer.Prove(ir::NE::Make(e1, e2)).has_value());
|
||||
|
||||
EXPECT_TRUE(analyzer.ProveGE(e1, e2).value() &&
|
||||
analyzer.Prove(e1 >= e2).value());
|
||||
EXPECT_FALSE(analyzer.ProveGE(e2, e1).has_value() ||
|
||||
analyzer.Prove(e2 >= e1).has_value());
|
||||
EXPECT_TRUE(analyzer.ProveLE(e2, e1).value() &&
|
||||
analyzer.Prove(e2 <= e1).value());
|
||||
EXPECT_FALSE(analyzer.ProveLE(e1, e2).has_value() ||
|
||||
analyzer.Prove(e1 <= e2).has_value());
|
||||
|
||||
EXPECT_FALSE(analyzer.ProveGT(e1, e2).has_value() ||
|
||||
analyzer.Prove(e1 > e2).has_value());
|
||||
EXPECT_FALSE(analyzer.ProveGT(e2, e1).value() &&
|
||||
analyzer.Prove(e2 > e1).value());
|
||||
EXPECT_FALSE(analyzer.ProveLT(e2, e1).has_value() ||
|
||||
analyzer.Prove(e2 < e1).has_value());
|
||||
EXPECT_FALSE(analyzer.ProveLT(e1, e2).value() &&
|
||||
analyzer.Prove(e1 < e2).value());
|
||||
|
||||
// case 2
|
||||
ir::Expr e3 = i + j + 1;
|
||||
ir::Expr e4 = i + j;
|
||||
|
||||
EXPECT_TRUE(analyzer.ProveEQ(e3, e3).value() &&
|
||||
analyzer.Prove(ir::EQ::Make(e3, e3)).value());
|
||||
EXPECT_FALSE(analyzer.ProveEQ(e3, e4).value() &&
|
||||
analyzer.Prove(ir::EQ::Make(e3, e4)).value());
|
||||
EXPECT_TRUE(analyzer.ProveNE(e3, e4).value() &&
|
||||
analyzer.Prove(ir::NE::Make(e3, e4)).value());
|
||||
EXPECT_FALSE(analyzer.ProveNE(e4, e4).value() &&
|
||||
analyzer.Prove(ir::NE::Make(e4, e4)).value());
|
||||
|
||||
EXPECT_TRUE(analyzer.ProveGE(e3, e4).value() &&
|
||||
analyzer.Prove(e3 >= e4).value());
|
||||
EXPECT_FALSE(analyzer.ProveGE(e4, e3).value() &&
|
||||
analyzer.Prove(e4 >= e3).value());
|
||||
EXPECT_TRUE(analyzer.ProveLE(e4, e3).value() &&
|
||||
analyzer.Prove(e4 <= e3).value());
|
||||
EXPECT_FALSE(analyzer.ProveLE(e3, e4).value() &&
|
||||
analyzer.Prove(e3 <= e4).value());
|
||||
|
||||
EXPECT_TRUE(analyzer.ProveGT(e3, e4).value() &&
|
||||
analyzer.Prove(e3 > e4).value());
|
||||
EXPECT_FALSE(analyzer.ProveGT(e4, e3).value() &&
|
||||
analyzer.Prove(e4 > e3).value());
|
||||
EXPECT_TRUE(analyzer.ProveLT(e4, e3).value() &&
|
||||
analyzer.Prove(e4 < e3).value());
|
||||
EXPECT_FALSE(analyzer.ProveLT(e3, e4).value() &&
|
||||
analyzer.Prove(e3 < e4).value());
|
||||
}
|
||||
|
||||
TEST_F(TestSymbolicExprAnalyzer, Divisible) {
|
||||
auto x = ir::Var(ir::Expr(1), ir::Expr(7), "x");
|
||||
auto y = ir::Var(ir::Expr(1), ir::Expr(15), "y");
|
||||
auto S = ir::Var(ir::Expr(16), ir::Expr(256), "S");
|
||||
|
||||
cas_intervals_t divisible_var_intervals = {
|
||||
{"x", CasInterval(x->lower_bound, x->upper_bound - ir::Expr(1))},
|
||||
{"y", CasInterval(y->lower_bound, y->upper_bound - ir::Expr(1))},
|
||||
{"S", CasInterval(S->lower_bound, S->upper_bound - ir::Expr(1))},
|
||||
};
|
||||
SymbolicExprAnalyzer divisible_analyzer{divisible_var_intervals};
|
||||
|
||||
// case 1
|
||||
ir::Expr e1 = 4 * x + 2 * y * x;
|
||||
ir::Expr e2 = x;
|
||||
ir::Expr e3 = y;
|
||||
|
||||
EXPECT_TRUE(divisible_analyzer.ProveDivisible(e1, e2).value_or(false));
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e1, e3).value_or(false));
|
||||
|
||||
// case 2
|
||||
ir::Expr e4 = y + y * x + 4 * y - x * y;
|
||||
|
||||
EXPECT_TRUE(divisible_analyzer.ProveDivisible(e4, e3).value_or(false));
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e4, e2).value_or(false));
|
||||
|
||||
// case 3
|
||||
ir::Expr e5 = x / y + x + y;
|
||||
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e5, e3).value_or(false));
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e5, e2).value_or(false));
|
||||
|
||||
// case 4
|
||||
ir::Expr e6 = S * x / 4 + x * y;
|
||||
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e6, e2).value_or(false));
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e6, e3).value_or(false));
|
||||
|
||||
ir::Expr e7 = 16 * x / 4 + x * y;
|
||||
|
||||
EXPECT_TRUE(divisible_analyzer.ProveDivisible(e7, e2).value_or(false));
|
||||
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e7, e3).value_or(false));
|
||||
}
|
||||
|
||||
TEST(SingleIntervalIntSet, constant) {
|
||||
SingleIntervalIntSet empty_set(ir::Expr(0), ir::Expr(-1));
|
||||
SingleIntervalIntSet all_set(SymbolicExprLimit::negative_inf,
|
||||
SymbolicExprLimit::positive_inf);
|
||||
SingleIntervalIntSet single_point(ir::Expr(0), ir::Expr(0));
|
||||
SingleIntervalIntSet interval_0_2_set(ir::Expr(0), ir::Expr(2));
|
||||
SingleIntervalIntSet interval_0_4_set(ir::Expr(0), ir::Expr(4));
|
||||
SingleIntervalIntSet interval_2_6_set(ir::Expr(2), ir::Expr(6));
|
||||
SingleIntervalIntSet interval_8_9_set(ir::Expr(8), ir::Expr(9));
|
||||
|
||||
EXPECT_TRUE(empty_set.ProveEmpty().value());
|
||||
EXPECT_FALSE(empty_set.ProveAll().value());
|
||||
EXPECT_FALSE(all_set.ProveEmpty().value());
|
||||
EXPECT_TRUE(all_set.ProveAll().value());
|
||||
EXPECT_TRUE(single_point.ProvePoint().value());
|
||||
EXPECT_FALSE(interval_0_2_set.ProvePoint().value());
|
||||
EXPECT_TRUE(interval_0_2_set.ProveSubSet(interval_0_4_set).value());
|
||||
EXPECT_FALSE(interval_0_4_set.ProveSubSet(interval_0_2_set).value());
|
||||
EXPECT_FALSE(interval_0_2_set.ProveSuperSet(interval_0_4_set).value());
|
||||
EXPECT_TRUE(interval_0_4_set.ProveSuperSet(interval_0_2_set).value());
|
||||
|
||||
EXPECT_TRUE(ProveEQ(interval_0_2_set, interval_0_2_set).value());
|
||||
EXPECT_FALSE(ProveEQ(interval_0_2_set, interval_0_4_set).value());
|
||||
|
||||
SingleIntervalIntSet union_0_6_set =
|
||||
ProvedUnion(interval_0_2_set, interval_2_6_set).value();
|
||||
EXPECT_EQ(union_0_6_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(union_0_6_set.Max(), ir::Expr(6));
|
||||
union_0_6_set = ProvedUnion(interval_2_6_set, interval_0_2_set).value();
|
||||
EXPECT_EQ(union_0_6_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(union_0_6_set.Max(), ir::Expr(6));
|
||||
SingleIntervalIntSet union_0_4_set =
|
||||
ProvedUnion(interval_0_2_set, interval_0_4_set).value();
|
||||
EXPECT_EQ(union_0_4_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(union_0_4_set.Max(), ir::Expr(4));
|
||||
union_0_4_set = ProvedUnion(interval_0_4_set, interval_0_2_set).value();
|
||||
EXPECT_EQ(union_0_4_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(union_0_4_set.Max(), ir::Expr(4));
|
||||
SingleIntervalIntSet union_0_9_set =
|
||||
ProvedUnion(interval_0_4_set, interval_8_9_set).value();
|
||||
EXPECT_EQ(union_0_9_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(union_0_9_set.Max(), ir::Expr(9));
|
||||
union_0_9_set = ProvedUnion(interval_8_9_set, interval_0_4_set).value();
|
||||
EXPECT_EQ(union_0_9_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(union_0_9_set.Max(), ir::Expr(9));
|
||||
|
||||
SingleIntervalIntSet intersect_0_2_set =
|
||||
ProvedIntersect(interval_0_2_set, interval_0_4_set).value();
|
||||
EXPECT_EQ(intersect_0_2_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(intersect_0_2_set.Max(), ir::Expr(2));
|
||||
intersect_0_2_set =
|
||||
ProvedIntersect(interval_0_4_set, interval_0_2_set).value();
|
||||
EXPECT_EQ(intersect_0_2_set.Min(), ir::Expr(0));
|
||||
EXPECT_EQ(intersect_0_2_set.Max(), ir::Expr(2));
|
||||
SingleIntervalIntSet intersect_2_2_set =
|
||||
ProvedIntersect(interval_0_2_set, interval_2_6_set).value();
|
||||
EXPECT_EQ(intersect_2_2_set.Min(), ir::Expr(2));
|
||||
EXPECT_EQ(intersect_2_2_set.Max(), ir::Expr(2));
|
||||
intersect_2_2_set =
|
||||
ProvedIntersect(interval_2_6_set, interval_0_2_set).value();
|
||||
EXPECT_EQ(intersect_2_2_set.Min(), ir::Expr(2));
|
||||
EXPECT_EQ(intersect_2_2_set.Max(), ir::Expr(2));
|
||||
SingleIntervalIntSet intersect_empty_set =
|
||||
ProvedIntersect(interval_0_4_set, interval_8_9_set).value();
|
||||
EXPECT_TRUE(intersect_empty_set.ProveEmpty().value());
|
||||
intersect_empty_set =
|
||||
ProvedIntersect(interval_8_9_set, interval_0_4_set).value();
|
||||
EXPECT_TRUE(intersect_empty_set.ProveEmpty().value());
|
||||
}
|
||||
|
||||
TEST(SingleIntervalIntSet, case_0) {
|
||||
ir::Var S0 = ir::Var(ir::Expr(0), ir::Expr(7), "S0");
|
||||
ir::Expr e1 = S0 * 16;
|
||||
ir::Expr e2 = S0 * 16 + 7;
|
||||
ir::Expr e3 = S0 * 16 + 15;
|
||||
SingleIntervalIntSet empty_set(e2, e1);
|
||||
SingleIntervalIntSet single_point(e3, e3);
|
||||
SingleIntervalIntSet set_0(e1, e2);
|
||||
SingleIntervalIntSet set_1(e1, e3);
|
||||
|
||||
EXPECT_TRUE(empty_set.ProveEmpty().value());
|
||||
EXPECT_FALSE(empty_set.ProveAll().value());
|
||||
EXPECT_TRUE(single_point.ProvePoint().value());
|
||||
EXPECT_FALSE(set_0.ProvePoint().value());
|
||||
EXPECT_TRUE(ProveEQ(set_0, set_0).value());
|
||||
EXPECT_FALSE(ProveEQ(set_0, set_1).value());
|
||||
|
||||
EXPECT_TRUE(set_0.ProveSubSet(set_1).value());
|
||||
EXPECT_FALSE(set_1.ProveSubSet(set_0).value());
|
||||
EXPECT_FALSE(set_0.ProveSuperSet(set_1).value());
|
||||
EXPECT_TRUE(set_1.ProveSuperSet(set_0).value());
|
||||
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_0, set_1).value(), set_1).value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedIntersect(set_0, set_1).value(), set_0).value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_1, single_point).value(), set_1).value());
|
||||
EXPECT_TRUE(
|
||||
ProveEQ(ProvedIntersect(set_1, single_point).value(), single_point)
|
||||
.value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_0, empty_set).value(), set_0).value());
|
||||
EXPECT_TRUE(
|
||||
ProveEQ(ProvedIntersect(set_0, empty_set).value(), empty_set).value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_0, single_point).value(), set_1).value());
|
||||
EXPECT_TRUE(ProvedIntersect(set_0, single_point).value().ProveEmpty());
|
||||
}
|
||||
|
||||
TEST(SingleIntervalIntSet, case_1) {
|
||||
ir::Var S0 = ir::Var(ir::Expr(0), ir::Expr(7), "S0");
|
||||
ir::Var S1 = ir::Var(ir::Expr(0), ir::Expr(15), "S1");
|
||||
ir::Expr e1 = S0 * 16;
|
||||
ir::Expr e2 = S0 * 16 + S1;
|
||||
ir::Expr e3 = S0 * 16 + S1 * 2 + 1;
|
||||
SingleIntervalIntSet empty_set(e3, e1);
|
||||
SingleIntervalIntSet single_point(e3, e3);
|
||||
SingleIntervalIntSet set_0(e1, e2);
|
||||
SingleIntervalIntSet set_1(e1, e3);
|
||||
|
||||
EXPECT_TRUE(empty_set.ProveEmpty().value());
|
||||
EXPECT_FALSE(empty_set.ProveAll().value());
|
||||
EXPECT_TRUE(single_point.ProvePoint().value());
|
||||
EXPECT_FALSE(set_0.ProvePoint().has_value());
|
||||
EXPECT_TRUE(ProveEQ(set_0, set_0).value());
|
||||
EXPECT_FALSE(ProveEQ(set_0, set_1).value());
|
||||
|
||||
EXPECT_TRUE(set_0.ProveSubSet(set_1).value());
|
||||
EXPECT_FALSE(set_1.ProveSubSet(set_0).value());
|
||||
EXPECT_FALSE(set_0.ProveSuperSet(set_1).value());
|
||||
EXPECT_TRUE(set_1.ProveSuperSet(set_0).value());
|
||||
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_0, set_1).value(), set_1).value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedIntersect(set_0, set_1).value(), set_0).value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_1, single_point).value(), set_1).value());
|
||||
EXPECT_TRUE(
|
||||
ProveEQ(ProvedIntersect(set_1, single_point).value(), single_point)
|
||||
.value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_0, empty_set).value(), set_0).value());
|
||||
EXPECT_TRUE(
|
||||
ProveEQ(ProvedIntersect(set_0, empty_set).value(), empty_set).value());
|
||||
EXPECT_TRUE(ProveEQ(ProvedUnion(set_0, single_point).value(), set_1).value());
|
||||
EXPECT_TRUE(
|
||||
ProvedIntersect(set_0, single_point).value().ProveEmpty().value());
|
||||
}
|
||||
|
||||
TEST(SingleIntervalIntSet, case_2) {
|
||||
ir::Var S = ir::Var(ir::Expr(0), ir::Expr(1), "S"); // S ∈ [0, 1)
|
||||
|
||||
SingleIntervalIntSet set_0{S, S + Expr(1)}; // [0, 1]
|
||||
SingleIntervalIntSet set_1{Expr(0), Expr(1)}; // [0, 1]
|
||||
SingleIntervalIntSet set_2{Expr(0), Expr(2)}; // [0, 2]
|
||||
|
||||
EXPECT_TRUE(ProveEQ(set_0, set_1).value());
|
||||
EXPECT_FALSE(ProveEQ(set_0, set_2).value());
|
||||
EXPECT_TRUE(set_0.ProveSubSet(set_2).value());
|
||||
EXPECT_TRUE(set_2.ProveSuperSet(set_0).value());
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,38 @@
|
||||
// Copyright (c) 2023 CINN 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 "paddle/cinn/common/is_reachable_predicator.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
TEST(IsReachablePredicator, simple) {
|
||||
IsReachablePredicator<int> IsReachable(
|
||||
// Get min depth
|
||||
[](int x) { return std::abs(x); },
|
||||
// Get max depth
|
||||
[](int x) { return std::abs(x); },
|
||||
// visit next node
|
||||
[](int x, const std::function<void(int)>& Handler) {
|
||||
Handler(x + (x / std::abs(x)));
|
||||
});
|
||||
EXPECT_TRUE(IsReachable(33, 99, [](int) {}));
|
||||
EXPECT_FALSE(IsReachable(33, -99, [](int) {}));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,55 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/common/shared.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/common/object.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
struct A : public Object {
|
||||
const char *type_info() const override { return "A"; }
|
||||
|
||||
Shared<A> other;
|
||||
};
|
||||
|
||||
class B : public Object {};
|
||||
|
||||
TEST(Shared, test) {
|
||||
Shared<A> a_ref(make_shared<A>());
|
||||
ASSERT_EQ(ref_count(a_ref.get()).val(), 1);
|
||||
|
||||
{ // local copy
|
||||
Shared<A> b = a_ref;
|
||||
EXPECT_EQ(ref_count(a_ref.get()).val(), 2);
|
||||
ASSERT_EQ(ref_count(b.get()).val(), 2);
|
||||
}
|
||||
|
||||
ASSERT_EQ(ref_count(a_ref.get()).val(), 1);
|
||||
}
|
||||
|
||||
TEST(Shared, cycle_share) {
|
||||
{
|
||||
Shared<A> a_ref(make_shared<A>());
|
||||
a_ref->other = a_ref;
|
||||
ASSERT_EQ(a_ref->__ref_count__.val(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,51 @@
|
||||
// Copyright (c) 2023 CINN 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 "paddle/cinn/common/topo_walker.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace cinn {
|
||||
namespace common {
|
||||
|
||||
TEST(TopoWalker, simple) {
|
||||
std::vector<std::pair<int, int>> edges{
|
||||
{0, 3}, {1, 2}, {1, 3}, {2, 3}, {3, 4}};
|
||||
TopoWalker<int> visitor(
|
||||
[&](int node, const std::function<void(int)>& NodeHandler) {
|
||||
for (const auto& pair : edges) {
|
||||
if (pair.second == node) {
|
||||
NodeHandler(pair.first);
|
||||
}
|
||||
}
|
||||
},
|
||||
[&](int node, const std::function<void(int)>& NodeHandler) {
|
||||
for (const auto& pair : edges) {
|
||||
if (pair.first == node) {
|
||||
NodeHandler(pair.second);
|
||||
}
|
||||
}
|
||||
});
|
||||
std::vector<int> sources{0, 1};
|
||||
std::vector<int> outputs;
|
||||
visitor(sources.begin(), sources.end(), [&](int node) {
|
||||
outputs.push_back(node);
|
||||
});
|
||||
std::vector<int> expected{0, 1, 2, 3, 4};
|
||||
EXPECT_TRUE((outputs == expected));
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,31 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/common/type.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace cinn::common {
|
||||
|
||||
TEST(Type, basic) {
|
||||
LOG(INFO) << I32();
|
||||
|
||||
auto i32 = I32();
|
||||
LOG(INFO) << I32();
|
||||
|
||||
LOG(INFO) << F32();
|
||||
LOG(INFO) << type_of<float>();
|
||||
}
|
||||
|
||||
} // namespace cinn::common
|
||||
@@ -0,0 +1 @@
|
||||
add_subdirectory(pe)
|
||||
@@ -0,0 +1 @@
|
||||
cinn_cc_test(test_load_params SRCS load_params_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,65 @@
|
||||
// Copyright (c) 2021 CINN 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 <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/hlir/pe/schedule.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace hlir {
|
||||
namespace pe {
|
||||
using ir::Tensor;
|
||||
|
||||
TEST(load_x86_params, load_x86_params) {
|
||||
auto &res = ScheduleParam::get_x86_instance().GetParam();
|
||||
std::string key =
|
||||
"X86ScheduleConv input 1 3 224 224 weight 64 3 7 7 stride 2 2 padding 3 "
|
||||
"3 dilation 1 1";
|
||||
ASSERT_EQ(res.count(key), 1);
|
||||
|
||||
paddle::flat_hash_map<std::string, int> conv2d_factors;
|
||||
auto target = cinn::common::DefaultHostTarget();
|
||||
std::vector<int> shape_input = {1, 64, 56, 56};
|
||||
std::vector<int> shape_weights = {64, 64, 3, 3};
|
||||
std::vector<int> strides = {1, 1};
|
||||
std::vector<int> pads = {1, 1};
|
||||
std::vector<int> dilations = {1, 1};
|
||||
key =
|
||||
GenerateX86ConvKey(shape_input, shape_weights, strides, pads, dilations);
|
||||
GetConv2dFactors(&conv2d_factors, -1, -1, -1, -1, -1, Float(32), target, key);
|
||||
int ic_bn_size = conv2d_factors["ic_bn"];
|
||||
int oc_bn_size = conv2d_factors["oc_bn"];
|
||||
int fc_bn_size = conv2d_factors["fc_bn"];
|
||||
int ow_bn_size = conv2d_factors["ow_bn"];
|
||||
int unroll_kw = conv2d_factors["unroll_kw"];
|
||||
ASSERT_EQ(ic_bn_size, 64);
|
||||
ASSERT_EQ(fc_bn_size, 64);
|
||||
ASSERT_EQ(oc_bn_size, 32);
|
||||
ASSERT_EQ(ow_bn_size, 7);
|
||||
ASSERT_EQ(unroll_kw, 1);
|
||||
}
|
||||
|
||||
TEST(load_cuda_params, load_cuda_params) {
|
||||
auto &res = ScheduleParam::get_cuda_instance().GetParam();
|
||||
if (res.empty()) {
|
||||
CreateCudaSerialData();
|
||||
LoadSerialData(&res);
|
||||
}
|
||||
std::string key = "CudaDirectConvSchedule 1 3 230 230 64 3 7 7 1 64 112 112";
|
||||
ASSERT_EQ(res.count(key), 1);
|
||||
}
|
||||
|
||||
} // namespace pe
|
||||
} // namespace hlir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1 @@
|
||||
add_subdirectory(test)
|
||||
@@ -0,0 +1,9 @@
|
||||
# cinn_cc_test(test_ir SRCS ir_test.cc DEPS core)
|
||||
# cinn_cc_test(test_ir_printer SRCS ir_printer_test.cc DEPS core)
|
||||
# cinn_cc_test(test_ir_operators SRCS ir_operators_test.cc DEPS core)
|
||||
# cinn_cc_test(test_tensor SRCS tensor_test.cc DEPS core)
|
||||
cinn_cc_test(test_collect_ir_nodes SRCS collect_ir_nodes_test.cc DEPS cinncore)
|
||||
|
||||
cinn_cc_test(test_intrinsic_ops SRCS intrinsic_ops_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_ir_verify SRCS ir_verify_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_ir_copy SRCS ir_copy_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,65 @@
|
||||
// Copyright (c) 2021 CINN 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 <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/cinn.h"
|
||||
#include "paddle/cinn/ir/ir.h"
|
||||
#include "paddle/cinn/ir/utils/stmt_converter.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace ir {
|
||||
namespace ir_utils {
|
||||
|
||||
TEST(CollectIRNodes, basic0) {
|
||||
Expr C = Expr(1) + 2;
|
||||
|
||||
auto exprs =
|
||||
CollectIRNodes(C, [](const Expr* x) { return x->As<ir::Add>(); });
|
||||
ASSERT_EQ(exprs.size(), 1UL);
|
||||
|
||||
auto ints =
|
||||
CollectIRNodes(C, [](const Expr* x) { return x->As<ir::IntImm>(); });
|
||||
ASSERT_EQ(ints.size(), 2UL);
|
||||
}
|
||||
|
||||
TEST(CollectIRNodes, basic) {
|
||||
Expr M(100);
|
||||
Expr N(200);
|
||||
Placeholder<float> A("A", {M, N});
|
||||
Placeholder<float> B("B", {M, N});
|
||||
|
||||
auto C = Compute(
|
||||
{M, N}, [&](Var i, Var j) { return A(i, j) + B(i, j); }, "C");
|
||||
|
||||
ast_gen_ius::TensorGroup tensor_group({C});
|
||||
|
||||
auto fn = LowerToAst("fn", {A, B, C}, &tensor_group);
|
||||
|
||||
LOG(INFO) << "fn:\n" << fn;
|
||||
|
||||
Expr expr_func_body = ConvertStmtBlockToExprBlock(fn->body_block);
|
||||
|
||||
auto tensors = CollectIRNodes(expr_func_body,
|
||||
[](const Expr* x) { return x->as_tensor(); });
|
||||
ASSERT_EQ(tensors.size(), 3UL);
|
||||
|
||||
LOG(INFO) << "fn.body:\n" << expr_func_body;
|
||||
auto tensors2 = CollectIRNodes(expr_func_body,
|
||||
[](const Expr* x) { return x->as_tensor(); });
|
||||
auto exprs = CollectIRNodes(expr_func_body, [](const Expr* x) { return x; });
|
||||
}
|
||||
} // namespace ir_utils
|
||||
} // namespace ir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,31 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/ir/intrinsic_ops.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace cinn::ir {
|
||||
|
||||
TEST(IntrinsicOp, basic) {
|
||||
Expr buffer(1);
|
||||
buffer->set_type(type_of<cinn_buffer_t*>());
|
||||
auto op = intrinsics::BufferGetDataHandle::Make(buffer);
|
||||
auto* ptr = op.As<IntrinsicOp>();
|
||||
ASSERT_TRUE(ptr);
|
||||
auto* obj = llvm::dyn_cast<intrinsics::BufferGetDataHandle>(ptr);
|
||||
ASSERT_TRUE(obj);
|
||||
}
|
||||
|
||||
} // namespace cinn::ir
|
||||
@@ -0,0 +1,32 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/ir/utils/ir_copy.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/ir/ir_printer.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace ir {
|
||||
namespace ir_utils {
|
||||
|
||||
TEST(IrCopy, basic) {
|
||||
Expr a(1.f);
|
||||
auto aa = IRCopy(a);
|
||||
LOG(INFO) << "aa " << aa;
|
||||
}
|
||||
} // namespace ir_utils
|
||||
} // namespace ir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/ir/op/ir_operators.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
namespace cinn {
|
||||
namespace ir {
|
||||
|
||||
TEST(ir_operators, test) {
|
||||
Expr a(1);
|
||||
Expr b = a + 1;
|
||||
}
|
||||
|
||||
} // namespace ir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,23 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/ir/ir_printer.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <sstream>
|
||||
|
||||
namespace cinn {
|
||||
namespace ir {} // namespace ir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,31 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/ir/ir.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/utils/string.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace ir {
|
||||
|
||||
TEST(Expr, basic) {
|
||||
Expr a(1);
|
||||
auto b = Expr(a);
|
||||
LOG(INFO) << b.as_int32();
|
||||
}
|
||||
|
||||
} // namespace ir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,31 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/ir/utils/ir_verify.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace ir {
|
||||
namespace ir_utils {
|
||||
TEST(IrVerify, basic) {
|
||||
Expr a(1);
|
||||
Expr b(1);
|
||||
IrVerify(a + b);
|
||||
}
|
||||
} // namespace ir_utils
|
||||
} // namespace ir
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,2 @@
|
||||
cinn_cc_test(test_compute SRCS compute_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_cinn_packed_func SRCS packed_func_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,39 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/lang/compute.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/cinn.h"
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
#include "paddle/cinn/ir/tensor.h"
|
||||
#include "paddle/cinn/lang/placeholder.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace lang {
|
||||
|
||||
TEST(Call, basic) {
|
||||
Expr M(100);
|
||||
|
||||
Placeholder<float> x("x", {M, Expr(10)});
|
||||
Placeholder<float> y("y", {M, Expr(10)});
|
||||
|
||||
std::vector<ReturnType> return_types(
|
||||
{{Float(32), std::vector<Expr>{{M, Expr(20)}}, "C"}});
|
||||
auto tensors = CallLowered("lowered_fun0", {Expr(x), Expr(y)}, return_types);
|
||||
}
|
||||
|
||||
} // namespace lang
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,96 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/lang/packed_func.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/ir/ir_printer.h"
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
#include "paddle/cinn/utils/string.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace lang {
|
||||
|
||||
TEST(Function, test) {
|
||||
PackedFunc::body_t func_body = [](Args args, RetValue* ret) {
|
||||
int a = args[0];
|
||||
int b = args[1];
|
||||
*ret = (a + b);
|
||||
};
|
||||
PackedFunc func(func_body);
|
||||
|
||||
int c = func(1, 2);
|
||||
LOG(INFO) << "c " << c;
|
||||
}
|
||||
|
||||
TEST(Function, test1) {
|
||||
PackedFunc::body_t body = [](Args args, RetValue* ret) {
|
||||
auto* msg = static_cast<const char*>(args[0]);
|
||||
(*ret) = msg;
|
||||
};
|
||||
|
||||
PackedFunc func(body);
|
||||
const char* msg = "hello world";
|
||||
char* c = func(msg);
|
||||
LOG(INFO) << static_cast<char*>(c);
|
||||
}
|
||||
|
||||
TEST(Function, Expr) {
|
||||
PackedFunc::body_t body = [](Args args, RetValue* ret) {
|
||||
Expr a = args[0];
|
||||
Expr b = args[1];
|
||||
|
||||
ASSERT_EQ(a->__ref_count__.val(), 4);
|
||||
ASSERT_EQ(b->__ref_count__.val(), 4);
|
||||
|
||||
Expr c = a + b;
|
||||
(*ret) = CINNValue(c);
|
||||
};
|
||||
|
||||
PackedFunc func(body);
|
||||
|
||||
Expr a(1);
|
||||
Expr b(2);
|
||||
ASSERT_EQ(a->__ref_count__.val(), 1);
|
||||
ASSERT_EQ(b->__ref_count__.val(), 1);
|
||||
|
||||
Expr ret = func(a, b);
|
||||
|
||||
ASSERT_EQ(utils::GetStreamCnt(ret), "(1 + 2)");
|
||||
}
|
||||
|
||||
TEST(Function, ReturnMultiValue) {
|
||||
PackedFunc::body_t body = [](Args args, RetValue* ret) {
|
||||
int a = args[0];
|
||||
int b = args[1];
|
||||
int c = a + b;
|
||||
int d = a - b;
|
||||
|
||||
*ret = cinn::common::CINNValuePack{
|
||||
{cinn::common::CINNValue(c), cinn::common::CINNValue(d)}};
|
||||
};
|
||||
|
||||
PackedFunc func(body);
|
||||
|
||||
cinn::common::CINNValuePack ret = func(1, 2);
|
||||
int c = ret[0];
|
||||
int d = ret[1];
|
||||
|
||||
EXPECT_EQ(c, 3);
|
||||
EXPECT_EQ(d, -1);
|
||||
}
|
||||
|
||||
} // namespace lang
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,3 @@
|
||||
cinn_cc_test(test_cast_simplify SRCS cast_simplify_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_replace_cross_thread_reduction SRCS
|
||||
replace_cross_thread_reduction_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,62 @@
|
||||
// Copyright (c) 2021 CINN 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 <gtest/gtest.h>
|
||||
|
||||
#include "paddle/cinn/ir/ir_printer.h"
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
#include "paddle/cinn/optim/ir_simplify.h"
|
||||
#include "paddle/cinn/utils/string.h"
|
||||
namespace cinn::optim {
|
||||
|
||||
TEST(CastSimplify, same_type) {
|
||||
Var n("n");
|
||||
Expr a = ir::Cast::Make(Int(32), n);
|
||||
LOG(INFO) << n->type();
|
||||
LOG(INFO) << a;
|
||||
SimplifyCast(&a);
|
||||
ASSERT_EQ(utils::GetStreamCnt(a), "n");
|
||||
}
|
||||
|
||||
TEST(CastSimplify, Imm_int) {
|
||||
Expr a = ir::Cast::Make(Int(64), Expr(1));
|
||||
Expr c = ir::Cast::Make(Int(32), a);
|
||||
LOG(INFO) << c;
|
||||
SimplifyCast(&c);
|
||||
LOG(INFO) << c;
|
||||
ASSERT_EQ(utils::GetStreamCnt(c), "1");
|
||||
ASSERT_EQ(c.type(), Int(32));
|
||||
}
|
||||
|
||||
TEST(CastSimplify, Imm_double) {
|
||||
Expr a = ir::Cast::Make(Float(64), Expr(2.33));
|
||||
Expr c = ir::Cast::Make(Int(32), a);
|
||||
LOG(INFO) << c;
|
||||
SimplifyCast(&c);
|
||||
LOG(INFO) << c;
|
||||
ASSERT_EQ(utils::GetStreamCnt(c), "2");
|
||||
ASSERT_EQ(c.type(), Int(32));
|
||||
}
|
||||
|
||||
TEST(CastSimplify, Imm_uint) {
|
||||
Expr a = ir::Cast::Make(UInt(64), Expr(1));
|
||||
Expr c = ir::Cast::Make(UInt(32), a);
|
||||
LOG(INFO) << c;
|
||||
SimplifyCast(&c);
|
||||
LOG(INFO) << c;
|
||||
ASSERT_EQ(utils::GetStreamCnt(c), "1");
|
||||
ASSERT_EQ(c.type(), UInt(32));
|
||||
}
|
||||
|
||||
} // namespace cinn::optim
|
||||
@@ -0,0 +1,102 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/optim/replace_cross_thread_reduction.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/cinn/cinn.h"
|
||||
#include "paddle/cinn/ir/ir.h"
|
||||
#include "paddle/cinn/ir/ir_printer.h"
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
#include "paddle/cinn/ir/schedule/ir_schedule.h"
|
||||
#include "paddle/cinn/ir/utils/stmt_converter.h"
|
||||
#include "paddle/cinn/utils/string.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace optim {
|
||||
|
||||
TEST(CrossThreadReductionReplacer, basic) {
|
||||
#ifdef CINN_WITH_CUDA
|
||||
Context::Global().ResetNameId();
|
||||
Placeholder<float> A("A", {Expr(64), Expr(128)});
|
||||
Target target = cinn::common::DefaultNVGPUTarget();
|
||||
Module::Builder builder("reduce_sum", target);
|
||||
Var reduce_j(128, "reduce_j");
|
||||
ir::Tensor B = Compute(
|
||||
{Expr(64)},
|
||||
[&](Var i) { return lang::ReduceSum(A(i, reduce_j), {reduce_j}); },
|
||||
"B");
|
||||
ast_gen_ius::TensorGroup tensor_group({A, B});
|
||||
auto func = lang::LowerToAst("reduce_sum", {A, B}, &tensor_group);
|
||||
VLOG(6) << "original func\n" << func;
|
||||
|
||||
ir::Expr expr_func_body = ir::ConvertStmtBlockToExprBlock(func->body_block);
|
||||
ir::ModuleExpr mod_expr({expr_func_body});
|
||||
ir::IRSchedule ir_sch(mod_expr);
|
||||
|
||||
ir_sch.Bind(ir_sch.GetLoops("B")[0], "blockIdx.x");
|
||||
ir_sch.Bind(ir_sch.GetLoops("B")[1], "threadIdx.x");
|
||||
|
||||
ir::Expr block = ir_sch.GetBlock("B");
|
||||
block.As<ir::ScheduleBlockRealize>()
|
||||
->schedule_block.As<ir::ScheduleBlock>()
|
||||
->reduce_method = ir::BlockReduceMethod();
|
||||
|
||||
ir::Expr func_body = ir_sch.GetModule().GetExprs()[0];
|
||||
std::vector<ir::Argument> args{
|
||||
ir::Argument(ir::Var("A"), ir::Argument::IO::kInput),
|
||||
ir::Argument(ir::Var("B"), ir::Argument::IO::kOutput)};
|
||||
auto new_func = ir::_LoweredFunc_::Make("test_func", args, func_body, {});
|
||||
VLOG(6) << "After Bind: " << new_func->body;
|
||||
|
||||
ReplaceCrossThreadReduction(new_func);
|
||||
VLOG(6) << "After ReplaceCrossThreadReduction: " << new_func->body;
|
||||
|
||||
EXPECT_EQ(utils::GetStreamCnt(new_func->body), utils::Trim(R"ROC({
|
||||
ScheduleBlock(root)
|
||||
{
|
||||
{
|
||||
thread_bind[blockIdx.x] for (i, 0, 64)
|
||||
{
|
||||
ScheduleBlock(B__reduce_init)
|
||||
{
|
||||
i0 = axis.bind(i)
|
||||
{
|
||||
B__reduce_init[i0] = 0.00000000f
|
||||
}
|
||||
}
|
||||
thread_bind[threadIdx.x] for (reduce_j, 0, 128)
|
||||
{
|
||||
ScheduleBlock(B)
|
||||
{
|
||||
i0_0, i1 = axis.bind(i, reduce_j)
|
||||
{
|
||||
B[i0_0] = cinn_block_reduce_sum_fp32(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
)ROC"));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace optim
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,2 @@
|
||||
cinn_cc_test(test_compute_at_transform SRCS compute_at_transform_test.cc DEPS
|
||||
cinncore)
|
||||
@@ -0,0 +1,58 @@
|
||||
// Copyright (c) 2021 CINN 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 <gtest/gtest.h>
|
||||
|
||||
namespace cinn {
|
||||
namespace poly {
|
||||
|
||||
TEST(ComputeAtTransform2, basic) {
|
||||
isl::ctx ctx(isl_ctx_alloc());
|
||||
isl::set pdomain(ctx, "{ p[i,j]: 0<=i,j<100 }");
|
||||
isl::map ptransform(ctx,
|
||||
"{ p[i,j]->p[t0,t1,t2]: t0=i%4 and t1=i/4 and t2=j }");
|
||||
isl::set cdomain(ctx, "{ c[i,j,k]: 0<=i,j,k<50 }");
|
||||
isl::map ctransform(
|
||||
ctx, "{ c[i,j,k]->c[t0,t1,t2,t3]: t0=i/4 and t1=i%4 and t2=j and t3=k }");
|
||||
|
||||
isl::map access(
|
||||
ctx, "{ c[i,j,k]->p[i,j]; c[i,j,k]->p[i+1,j]; c[i,j,k]->p[i-1,j] }");
|
||||
|
||||
poly::ComputeAtTransform t(
|
||||
pdomain, cdomain, access, ptransform, ctransform, 1);
|
||||
t();
|
||||
|
||||
t.DisplayC();
|
||||
|
||||
isl::map pschedule(ctx,
|
||||
"{ p[i0,i1,i2,i3,i4] -> [t0,t1,t1t, t2,t3,t4,t5]: t0=i0 "
|
||||
"and t1=i1 and t2=i2 and t3=i3 and t4=i4 "
|
||||
"and t5=0 and t1t=0 }");
|
||||
isl::map cschedule(
|
||||
ctx,
|
||||
"[_c_0,_c_1] -> { c[i0,i1,i2,i3] -> [t0,t1,t1t,t2,t3,t4,t5]: t0=i0 and "
|
||||
"t1=i1 and t2=i2 and t3=i3 "
|
||||
"and t4=0 and t5=0 and t1t=1 }");
|
||||
|
||||
t.DisplayC(pschedule.release(), cschedule.release());
|
||||
|
||||
LOG(INFO) << "shape:";
|
||||
auto shape = t.GetProducerAdjustedShape();
|
||||
for (int i = 0; i < shape.size(); i++) {
|
||||
LOG(INFO) << shape[i];
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace poly
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,2 @@
|
||||
cinn_cc_test(test_cinn_runtime SRCS cinn_runtime_test.cc DEPS cinn_runtime)
|
||||
add_subdirectory(cuda)
|
||||
@@ -0,0 +1,50 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/runtime/cinn_runtime.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
TEST(buffer, basic) {
|
||||
auto* buffer =
|
||||
cinn_buffer_t::new_(cinn_x86_device, cinn_float32_t(), {3, 10});
|
||||
ASSERT_TRUE(buffer);
|
||||
ASSERT_TRUE(buffer->device_interface);
|
||||
ASSERT_EQ(buffer->device_interface, cinn_x86_device_interface());
|
||||
buffer->device_interface->impl->malloc(NULL, buffer);
|
||||
auto* data = reinterpret_cast<float*>(buffer->memory);
|
||||
data[0] = 0.f;
|
||||
data[1] = 1.f;
|
||||
EXPECT_EQ(data[0], 0.f);
|
||||
EXPECT_EQ(data[1], 1.f);
|
||||
}
|
||||
|
||||
TEST(cinn_print_debug_string, basic) {
|
||||
cinn_print_debug_string("hello world");
|
||||
cinn_print_debug_string("should be 1, %d", 1);
|
||||
int a = 1;
|
||||
cinn_print_debug_string("should be pointer, %p", &a);
|
||||
cinn_print_debug_string("should be 1, %d", a);
|
||||
cinn_print_debug_string("v3[%d %d %d], ", 1, 2, 3);
|
||||
}
|
||||
|
||||
TEST(cinn_args_construct, basic) {
|
||||
cinn_pod_value_t arr[4];
|
||||
cinn_pod_value_t a0(0);
|
||||
cinn_pod_value_t a1(1);
|
||||
cinn_pod_value_t a2(2);
|
||||
cinn_pod_value_t a3(3);
|
||||
cinn_args_construct(arr, 4, &a0, &a1, &a2, &a3);
|
||||
for (int i = 0; i < 4; i++) ASSERT_EQ((int)arr[i], i);
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
if(NOT WITH_CUDA)
|
||||
return()
|
||||
endif()
|
||||
|
||||
cinn_nv_test(test_cuda_module SRCS cuda_module_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,205 @@
|
||||
// Copyright (c) 2021 CINN 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 "paddle/cinn/runtime/cuda/cuda_module.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <random>
|
||||
|
||||
#include "paddle/cinn/backends/nvrtc/nvrtc_util.h"
|
||||
#include "paddle/cinn/cinn.h"
|
||||
#include "paddle/cinn/runtime/cuda/cuda_util.h"
|
||||
#include "paddle/cinn/runtime/cuda/test_util.h"
|
||||
#include "paddle/cinn/runtime/cuda/use_extern_funcs.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace runtime {
|
||||
namespace cuda {
|
||||
|
||||
TEST(CUDAModule, basic) {
|
||||
backends::nvrtc::Compiler compiler;
|
||||
|
||||
std::string source_code = R"ROC(
|
||||
extern "C" __global__
|
||||
void saxpy(float a, float *x, float *y, float *out, size_t n)
|
||||
{
|
||||
size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (tid < n) {
|
||||
out[tid] = a * x[tid] + y[tid];
|
||||
}
|
||||
}
|
||||
)ROC";
|
||||
|
||||
auto ptx = compiler(source_code);
|
||||
PADDLE_ENFORCE_NE(
|
||||
ptx.empty(), true, ::common::errors::NotFound("ptx is empty!"));
|
||||
|
||||
CUDAModule module(ptx, CUDAModule::Kind::PTX);
|
||||
auto func = module.GetFunction(0, "saxpy");
|
||||
ASSERT_TRUE(func);
|
||||
}
|
||||
|
||||
TEST(CUDAModule, float16) {
|
||||
using cinn::common::float16;
|
||||
using runtime::cuda::util::Vector;
|
||||
|
||||
auto generate_ptx = [] {
|
||||
backends::nvrtc::Compiler compiler;
|
||||
|
||||
std::string source_code = R"(
|
||||
#include <cstdint>
|
||||
#define CINN_WITH_CUDA
|
||||
#include "float16.h"
|
||||
using cinn::common::float16;
|
||||
|
||||
extern "C" __global__
|
||||
void cast_fp32_to_fp16_cuda_kernel(const float* input, const int num, float16* output) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
output[idx] = float16(input[idx]);
|
||||
}
|
||||
}
|
||||
)";
|
||||
|
||||
auto ptx = compiler(source_code);
|
||||
PADDLE_ENFORCE_NE(
|
||||
ptx.empty(), true, ::common::errors::NotFound("ptx is empty!"));
|
||||
return ptx;
|
||||
};
|
||||
|
||||
auto ptx = generate_ptx();
|
||||
|
||||
CUDAModule cuda_module(ptx, CUDAModule::Kind::PTX);
|
||||
auto func = cuda_module.GetFunction(0, "cast_fp32_to_fp16_cuda_kernel");
|
||||
ASSERT_TRUE(func);
|
||||
|
||||
int size = 100;
|
||||
dim3 blocks_per_grid(1);
|
||||
dim3 threads_per_block(100);
|
||||
|
||||
std::vector<float> x_host(size);
|
||||
{
|
||||
std::random_device r;
|
||||
std::default_random_engine eng(r());
|
||||
std::uniform_real_distribution<float> dis(1e-5f, 1.0f);
|
||||
for (size_t i = 0; i < x_host.size(); ++i) {
|
||||
x_host[i] = dis(eng);
|
||||
}
|
||||
}
|
||||
Vector<float> x_device(x_host);
|
||||
Vector<float16> y_device(size);
|
||||
auto* x_p{x_device.data()};
|
||||
auto* y_p{y_device.data()};
|
||||
|
||||
void* args[] = {&x_p, &size, &y_p};
|
||||
cuda_module.LaunchKernel(0,
|
||||
"cast_fp32_to_fp16_cuda_kernel",
|
||||
blocks_per_grid,
|
||||
threads_per_block,
|
||||
args);
|
||||
CUDA_CALL(cudaDeviceSynchronize());
|
||||
|
||||
std::vector<float16> y_host = y_device.to_host();
|
||||
bool res = std::equal(x_host.begin(),
|
||||
x_host.end(),
|
||||
y_host.begin(),
|
||||
[](float x, float16 y) -> bool {
|
||||
return std::abs(x - static_cast<float>(y)) < 1e-2f;
|
||||
});
|
||||
PADDLE_ENFORCE_EQ(
|
||||
res,
|
||||
true,
|
||||
::common::errors::PreconditionNotMet(
|
||||
"The difference between two arrays exceeds the bound."));
|
||||
}
|
||||
|
||||
TEST(CUDAModule, bfloat16) {
|
||||
using cinn::common::bfloat16;
|
||||
using runtime::cuda::util::Vector;
|
||||
|
||||
auto generate_ptx = [] {
|
||||
backends::nvrtc::Compiler compiler;
|
||||
|
||||
std::string source_code = R"(
|
||||
#include <cstdint>
|
||||
#define CINN_WITH_CUDA
|
||||
#include "bfloat16.h"
|
||||
using cinn::common::bfloat16;
|
||||
|
||||
extern "C" __global__
|
||||
void cast_fp32_to_bf16_cuda_kernel(const float* input, const int num, bfloat16* output) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num) {
|
||||
output[idx] = bfloat16(input[idx]);
|
||||
}
|
||||
}
|
||||
)";
|
||||
|
||||
auto ptx = compiler(source_code);
|
||||
PADDLE_ENFORCE_NE(
|
||||
ptx.empty(), true, ::common::errors::NotFound("ptx is empty!"));
|
||||
return ptx;
|
||||
};
|
||||
|
||||
auto ptx = generate_ptx();
|
||||
|
||||
CUDAModule cuda_module(ptx, CUDAModule::Kind::PTX);
|
||||
auto func = cuda_module.GetFunction(0, "cast_fp32_to_bf16_cuda_kernel");
|
||||
ASSERT_TRUE(func);
|
||||
|
||||
int size = 100;
|
||||
dim3 blocks_per_grid(1);
|
||||
dim3 threads_per_block(100);
|
||||
|
||||
std::vector<float> x_host(size);
|
||||
{
|
||||
std::random_device r;
|
||||
std::default_random_engine eng(r());
|
||||
std::uniform_real_distribution<float> dis(1e-5f, 1.0f);
|
||||
for (size_t i = 0; i < x_host.size(); ++i) {
|
||||
x_host[i] = dis(eng);
|
||||
}
|
||||
}
|
||||
Vector<float> x_device(x_host);
|
||||
Vector<bfloat16> y_device(size);
|
||||
auto* x_p{x_device.data()};
|
||||
auto* y_p{y_device.data()};
|
||||
|
||||
void* args[] = {&x_p, &size, &y_p};
|
||||
cuda_module.LaunchKernel(0,
|
||||
"cast_fp32_to_bf16_cuda_kernel",
|
||||
blocks_per_grid,
|
||||
threads_per_block,
|
||||
args);
|
||||
CUDA_CALL(cudaDeviceSynchronize());
|
||||
|
||||
std::vector<bfloat16> y_host = y_device.to_host();
|
||||
bool res = std::equal(x_host.begin(),
|
||||
x_host.end(),
|
||||
y_host.begin(),
|
||||
[](float x, bfloat16 y) -> bool {
|
||||
return std::abs(x - static_cast<float>(y)) < 1e-2f;
|
||||
});
|
||||
PADDLE_ENFORCE_EQ(
|
||||
res,
|
||||
true,
|
||||
::common::errors::PreconditionNotMet(
|
||||
"The difference between two arrays exceeds the bound."));
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace runtime
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,2 @@
|
||||
cinn_cc_test(test_sized_multi_set SRCS sized_multi_set_test.cc DEPS cinncore)
|
||||
cinn_cc_test(test_multi_threading SRCS multi_threading_test.cc DEPS cinncore)
|
||||
@@ -0,0 +1,64 @@
|
||||
// Copyright (c) 2022 CINN 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 "paddle/cinn/utils/multi_threading.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/common/enforce.h"
|
||||
|
||||
namespace cinn {
|
||||
namespace utils {
|
||||
|
||||
TEST(JobDispatcher, SequenceDispatcher) {
|
||||
std::unique_ptr<JobDispatcher> dispatcher =
|
||||
std::make_unique<SequenceDispatcher>(1, 3);
|
||||
ASSERT_EQ(1, dispatcher->Next());
|
||||
ASSERT_EQ(2, dispatcher->Next());
|
||||
// check reach the end
|
||||
ASSERT_EQ(-1, dispatcher->Next());
|
||||
}
|
||||
|
||||
TEST(parallel_run, Basic) {
|
||||
std::vector<int> results(100, -1);
|
||||
auto worker_fn = [&results](int index) {
|
||||
PADDLE_ENFORCE_LT(index,
|
||||
results.size(),
|
||||
::common::errors::InvalidArgument("invalid index!"));
|
||||
results[index] = index;
|
||||
};
|
||||
// check process every index in the extent of [0, 100) with step 1
|
||||
parallel_run(worker_fn, SequenceDispatcher(0, 100), 2);
|
||||
for (int i = 0; i < 100; ++i) {
|
||||
ASSERT_EQ(results[i], i);
|
||||
}
|
||||
|
||||
// check only indexes in the extent of [0, 100) with step 3 are processed
|
||||
results.assign(100, -1);
|
||||
parallel_run(worker_fn, SequenceDispatcher(0, 100, 3), 3);
|
||||
for (int i = 0; i < 100; ++i) {
|
||||
if (i % 3 == 0) {
|
||||
ASSERT_EQ(results[i], i);
|
||||
} else {
|
||||
ASSERT_EQ(results[i], -1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace utils
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,82 @@
|
||||
// Copyright (c) 2022 CINN 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 "paddle/cinn/utils/sized_multi_set.h"
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace cinn {
|
||||
namespace utils {
|
||||
|
||||
TEST(SizedMultiSet, PopMax) {
|
||||
SizedMultiSet<int> sized_multi_set(5);
|
||||
|
||||
for (int i = 0; i < 10; ++i) {
|
||||
sized_multi_set.Push(i);
|
||||
if (i < 5) {
|
||||
EXPECT_EQ(sized_multi_set.Size(), static_cast<size_t>(i + 1));
|
||||
EXPECT_EQ(sized_multi_set.MaxValue(), i);
|
||||
EXPECT_EQ(sized_multi_set.MinValue(), 0);
|
||||
} else {
|
||||
EXPECT_EQ(sized_multi_set.Size(), 5);
|
||||
EXPECT_EQ(sized_multi_set.MaxValue(), 4);
|
||||
EXPECT_EQ(sized_multi_set.MinValue(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> vec = sized_multi_set.ReturnAsContainer<std::vector<int>>();
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
EXPECT_EQ(vec[i], i);
|
||||
}
|
||||
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
sized_multi_set.Pop();
|
||||
EXPECT_EQ(sized_multi_set.Size(), static_cast<size_t>(4 - i));
|
||||
EXPECT_EQ(sized_multi_set.MaxValue(), static_cast<size_t>(3 - i));
|
||||
EXPECT_EQ(sized_multi_set.MinValue(), static_cast<size_t>(0));
|
||||
}
|
||||
}
|
||||
|
||||
TEST(SizedMultiSet, PopMin) {
|
||||
SizedMultiSet<int> sized_multi_set(5, /* pop_max_when_full = */ false);
|
||||
for (int i = 0; i < 10; ++i) {
|
||||
sized_multi_set.Push(i);
|
||||
if (i < 5) {
|
||||
EXPECT_EQ(sized_multi_set.Size(), static_cast<size_t>(i + 1));
|
||||
EXPECT_EQ(sized_multi_set.MaxValue(), i);
|
||||
EXPECT_EQ(sized_multi_set.MinValue(), 0);
|
||||
} else {
|
||||
EXPECT_EQ(sized_multi_set.Size(), 5);
|
||||
EXPECT_EQ(sized_multi_set.MaxValue(), i);
|
||||
EXPECT_EQ(sized_multi_set.MinValue(), i - 4);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> vec = sized_multi_set.ReturnAsContainer<std::vector<int>>();
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
EXPECT_EQ(vec[i], i + 5);
|
||||
}
|
||||
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
sized_multi_set.Pop();
|
||||
EXPECT_EQ(sized_multi_set.Size(), static_cast<size_t>(4 - i));
|
||||
EXPECT_EQ(sized_multi_set.MaxValue(), static_cast<size_t>(9));
|
||||
EXPECT_EQ(sized_multi_set.MinValue(), static_cast<size_t>(6 + i));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace utils
|
||||
} // namespace cinn
|
||||
@@ -0,0 +1,248 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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.
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
|
||||
#include <ATen/cuda/CUDABlas.h>
|
||||
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/common/bfloat16.h"
|
||||
#include "paddle/phi/common/complex.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
|
||||
// Helper: allocate three same-sized device buffers, copy host data in,
|
||||
// invoke a kernel via |fn|, copy results back, synchronize, then free.
|
||||
// |fn| receives (d_a, d_b, d_c); it must not free them.
|
||||
template <typename T, typename Fn>
|
||||
static void runOnDevice(const std::vector<T>& h_a,
|
||||
const std::vector<T>& h_b,
|
||||
std::vector<T>* h_c,
|
||||
Fn fn) {
|
||||
size_t bytes = h_a.size() * sizeof(T);
|
||||
T *d_a = nullptr, *d_b = nullptr, *d_c = nullptr;
|
||||
|
||||
ASSERT_EQ(cudaMalloc(&d_a, bytes), cudaSuccess);
|
||||
ASSERT_EQ(cudaMalloc(&d_b, bytes), cudaSuccess);
|
||||
ASSERT_EQ(cudaMalloc(&d_c, bytes), cudaSuccess);
|
||||
|
||||
ASSERT_EQ(cudaMemcpy(d_a, h_a.data(), bytes, cudaMemcpyHostToDevice),
|
||||
cudaSuccess);
|
||||
ASSERT_EQ(cudaMemcpy(d_b, h_b.data(), bytes, cudaMemcpyHostToDevice),
|
||||
cudaSuccess);
|
||||
ASSERT_EQ(cudaMemcpy(d_c, h_c->data(), bytes, cudaMemcpyHostToDevice),
|
||||
cudaSuccess);
|
||||
|
||||
fn(d_a, d_b, d_c);
|
||||
|
||||
ASSERT_EQ(cudaMemcpy(h_c->data(), d_c, bytes, cudaMemcpyDeviceToHost),
|
||||
cudaSuccess);
|
||||
ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
|
||||
|
||||
cudaFree(d_a);
|
||||
cudaFree(d_b);
|
||||
cudaFree(d_c);
|
||||
}
|
||||
|
||||
// Runs 2x2 no-transpose gemm: C = alpha*A*B + beta*C and checks the result.
|
||||
//
|
||||
// Column-major layout:
|
||||
// A: col0={1,3}, col1={2,4} => logical A = [[1,2],[3,4]]
|
||||
// B: col0={5,7}, col1={6,8} => logical B = [[5,6],[7,8]]
|
||||
// A*B = [[19,22],[43,50]] stored col-major: col0={19,43}, col1={22,50}
|
||||
template <typename T, typename MathT = at::opmath_type<T>>
|
||||
class GemmTester {
|
||||
public:
|
||||
static constexpr int64_t N = 2;
|
||||
|
||||
static double toDouble(T val) { return static_cast<double>(val); }
|
||||
|
||||
void Run() {
|
||||
std::vector<T> h_a = {T(1), T(3), T(2), T(4)};
|
||||
std::vector<T> h_b = {T(5), T(7), T(6), T(8)};
|
||||
std::vector<T> h_c(N * N, T(0));
|
||||
|
||||
MathT alpha = static_cast<MathT>(1);
|
||||
MathT beta = static_cast<MathT>(0);
|
||||
|
||||
runOnDevice(h_a, h_b, &h_c, [&](T* d_a, T* d_b, T* d_c) {
|
||||
at::cuda::blas::gemm<T>(
|
||||
'N', 'N', N, N, N, alpha, d_a, N, d_b, N, beta, d_c, N);
|
||||
});
|
||||
|
||||
EXPECT_NEAR(toDouble(h_c[0]), 19.0, 1e-2); // C(0,0)
|
||||
EXPECT_NEAR(toDouble(h_c[1]), 43.0, 1e-2); // C(1,0)
|
||||
EXPECT_NEAR(toDouble(h_c[2]), 22.0, 1e-2); // C(0,1)
|
||||
EXPECT_NEAR(toDouble(h_c[3]), 50.0, 1e-2); // C(1,1)
|
||||
}
|
||||
|
||||
// transA='T': C = alpha * A^T * B + beta * C
|
||||
// A^T = [[1,3],[2,4]], A^T * B = [[26,30],[38,44]]
|
||||
void RunTransA() {
|
||||
std::vector<T> h_a = {T(1), T(3), T(2), T(4)};
|
||||
std::vector<T> h_b = {T(5), T(7), T(6), T(8)};
|
||||
std::vector<T> h_c(N * N, T(0));
|
||||
|
||||
MathT alpha = static_cast<MathT>(1);
|
||||
MathT beta = static_cast<MathT>(0);
|
||||
|
||||
runOnDevice(h_a, h_b, &h_c, [&](T* d_a, T* d_b, T* d_c) {
|
||||
at::cuda::blas::gemm<T>(
|
||||
'T', 'N', N, N, N, alpha, d_a, N, d_b, N, beta, d_c, N);
|
||||
});
|
||||
|
||||
EXPECT_NEAR(toDouble(h_c[0]), 26.0, 1e-2);
|
||||
EXPECT_NEAR(toDouble(h_c[1]), 38.0, 1e-2);
|
||||
EXPECT_NEAR(toDouble(h_c[2]), 30.0, 1e-2);
|
||||
EXPECT_NEAR(toDouble(h_c[3]), 44.0, 1e-2);
|
||||
}
|
||||
};
|
||||
|
||||
TEST(CUDABlasTest, GemmDouble) {
|
||||
GemmTester<double> t;
|
||||
t.Run();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmDoubleTransA) {
|
||||
GemmTester<double> t;
|
||||
t.RunTransA();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmFloat) {
|
||||
GemmTester<float> t;
|
||||
t.Run();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmFloatTransA) {
|
||||
GemmTester<float> t;
|
||||
t.RunTransA();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmFloatTransALowercase) {
|
||||
constexpr int64_t N = 2;
|
||||
|
||||
std::vector<float> h_a = {1.F, 3.F, 2.F, 4.F};
|
||||
std::vector<float> h_b = {5.F, 7.F, 6.F, 8.F};
|
||||
std::vector<float> h_c(N * N, 0.F);
|
||||
|
||||
float alpha = 1.F;
|
||||
float beta = 0.F;
|
||||
runOnDevice(h_a, h_b, &h_c, [&](float* d_a, float* d_b, float* d_c) {
|
||||
at::cuda::blas::gemm<float>(
|
||||
't', 'n', N, N, N, alpha, d_a, N, d_b, N, beta, d_c, N);
|
||||
});
|
||||
|
||||
EXPECT_NEAR(h_c[0], 26.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[1], 38.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[2], 30.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[3], 44.0f, 1e-3f);
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmComplexDouble) {
|
||||
GemmTester<c10::complex<double>> t;
|
||||
t.Run();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmComplexFloat) {
|
||||
GemmTester<c10::complex<float>> t;
|
||||
t.Run();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmHalf) {
|
||||
GemmTester<at::Half> t;
|
||||
t.Run();
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmBFloat16) {
|
||||
GemmTester<at::BFloat16> t;
|
||||
t.Run();
|
||||
}
|
||||
|
||||
// to_cublas_op 'C'/'c' path: C = A^H * I = A^H (conjugate-transpose of A).
|
||||
//
|
||||
// A stored col-major: col0={1+i,2+2i}, col1={3+3i,4+4i}
|
||||
// A^H stored col-major: col0={1-i,3-3i}, col1={2-2i,4-4i}
|
||||
TEST(CUDABlasTest, GemmComplexFloatConjTrans) {
|
||||
constexpr int64_t N = 2;
|
||||
using T = c10::complex<float>;
|
||||
|
||||
std::vector<T> h_a = {T(1, 1), T(2, 2), T(3, 3), T(4, 4)};
|
||||
std::vector<T> h_b = {T(1, 0), T(0, 0), T(0, 0), T(1, 0)}; // identity
|
||||
std::vector<T> h_c(N * N, T(0, 0));
|
||||
|
||||
float alpha = 1.0f;
|
||||
float beta = 0.0f;
|
||||
|
||||
runOnDevice(h_a, h_b, &h_c, [&](T* d_a, T* d_b, T* d_c) {
|
||||
at::cuda::blas::gemm<T>(
|
||||
'C', 'N', N, N, N, alpha, d_a, N, d_b, N, beta, d_c, N);
|
||||
});
|
||||
|
||||
EXPECT_NEAR(h_c[0].real, 1.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[0].imag, -1.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[1].real, 3.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[1].imag, -3.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[2].real, 2.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[2].imag, -2.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[3].real, 4.0f, 1e-3f);
|
||||
EXPECT_NEAR(h_c[3].imag, -4.0f, 1e-3f);
|
||||
}
|
||||
|
||||
// Same as above but uses lowercase 'c'/'n' to exercise that switch-case branch.
|
||||
TEST(CUDABlasTest, GemmComplexDoubleConjTransLower) {
|
||||
constexpr int64_t N = 2;
|
||||
using T = c10::complex<double>;
|
||||
|
||||
std::vector<T> h_a = {T(1, 1), T(2, 2), T(3, 3), T(4, 4)};
|
||||
std::vector<T> h_b = {T(1, 0), T(0, 0), T(0, 0), T(1, 0)};
|
||||
std::vector<T> h_c(N * N, T(0, 0));
|
||||
|
||||
double alpha = 1.0;
|
||||
double beta = 0.0;
|
||||
|
||||
runOnDevice(h_a, h_b, &h_c, [&](T* d_a, T* d_b, T* d_c) {
|
||||
at::cuda::blas::gemm<T>(
|
||||
'c', 'n', N, N, N, alpha, d_a, N, d_b, N, beta, d_c, N);
|
||||
});
|
||||
|
||||
EXPECT_NEAR(h_c[0].real, 1.0, 1e-6);
|
||||
EXPECT_NEAR(h_c[0].imag, -1.0, 1e-6);
|
||||
EXPECT_NEAR(h_c[1].real, 3.0, 1e-6);
|
||||
EXPECT_NEAR(h_c[1].imag, -3.0, 1e-6);
|
||||
}
|
||||
|
||||
TEST(CUDABlasTest, GemmInvalidTransposeThrows) {
|
||||
constexpr int64_t N = 1;
|
||||
double alpha = 1.0;
|
||||
double beta = 0.0;
|
||||
EXPECT_THROW(at::cuda::blas::gemm<double>('X',
|
||||
'N',
|
||||
N,
|
||||
N,
|
||||
N,
|
||||
alpha,
|
||||
static_cast<const double*>(nullptr),
|
||||
N,
|
||||
static_cast<const double*>(nullptr),
|
||||
N,
|
||||
beta,
|
||||
static_cast<double*>(nullptr),
|
||||
N),
|
||||
std::exception);
|
||||
}
|
||||
|
||||
#endif // PADDLE_WITH_CUDA
|
||||
@@ -0,0 +1,468 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/core/Allocator.h>
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <torch/cuda.h>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#endif
|
||||
|
||||
// Platform-specific definitions for memory operations
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
#include <hip/hip_runtime.h>
|
||||
#define MEMCPY_FN hipMemcpy
|
||||
#define MEMCPY_HOST_TO_DEVICE hipMemcpyHostToDevice
|
||||
#define MEMCPY_DEVICE_TO_HOST hipMemcpyDeviceToHost
|
||||
#define SUCCESS_CODE hipSuccess
|
||||
#define DEVICE_SYNCHRONIZE_FN hipDeviceSynchronize
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
#define MEMCPY_FN cudaMemcpy
|
||||
#define MEMCPY_HOST_TO_DEVICE cudaMemcpyHostToDevice
|
||||
#define MEMCPY_DEVICE_TO_HOST cudaMemcpyDeviceToHost
|
||||
#define SUCCESS_CODE cudaSuccess
|
||||
#define DEVICE_SYNCHRONIZE_FN cudaDeviceSynchronize
|
||||
#endif
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CUDAFunctions.h — covers the 2 missing lines:
|
||||
// c10::cuda::device_synchronize() and c10::cuda::stream_synchronize()
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
TEST(CUDAFunctionsTest, DeviceSynchronize) {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
// Exercises the PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize()) branch
|
||||
ASSERT_NO_THROW(c10::cuda::device_synchronize());
|
||||
#else
|
||||
// In CPU-only builds, device_synchronize throws
|
||||
ASSERT_THROW(c10::cuda::device_synchronize(), std::exception);
|
||||
#endif
|
||||
}
|
||||
|
||||
// CPU-only: torch::cuda::synchronize must report "No CUDA GPUs are available"
|
||||
// rather than the older "Cannot visit device count" produced by device_count().
|
||||
// Matches PyTorch behavior where device_count() returns 0 in CPU-only builds
|
||||
// and the synchronize() pre-check is the single source of the GPU-missing
|
||||
// error message.
|
||||
TEST(CUDAFunctionsTest, SynchronizeReportsNoGpuMessageInCpuOnly) {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// Only relevant in CPU-only builds
|
||||
return;
|
||||
#else
|
||||
try {
|
||||
torch::cuda::synchronize();
|
||||
FAIL() << "expected exception";
|
||||
} catch (const std::exception& e) {
|
||||
const std::string msg = e.what();
|
||||
EXPECT_NE(msg.find("No CUDA GPUs are available"), std::string::npos) << msg;
|
||||
EXPECT_EQ(msg.find("Cannot visit device count"), std::string::npos) << msg;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(CUDAFunctionsTest, StreamSynchronize) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
// Exercises phi::backends::gpu::GpuStreamSync()
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
ASSERT_NO_THROW(c10::cuda::stream_synchronize(stream));
|
||||
}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(CUDAFunctionsTest, AtNamespaceAliases) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
// Exercises the using aliases in at::cuda namespace
|
||||
ASSERT_NO_THROW(at::cuda::device_synchronize());
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
ASSERT_NO_THROW(at::cuda::stream_synchronize(stream));
|
||||
}
|
||||
|
||||
TEST(CUDAFunctionsTest, TorchSynchronizePreservesCurrentDevice) {
|
||||
if (!torch::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
if (torch::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
|
||||
constexpr int current_device = 0;
|
||||
constexpr int other_device = 1;
|
||||
c10::cuda::CUDAGuard guard(static_cast<c10::DeviceIndex>(current_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), current_device);
|
||||
|
||||
ASSERT_NO_THROW(torch::cuda::synchronize(other_device));
|
||||
EXPECT_EQ(phi::backends::gpu::GetCurrentDeviceId(), current_device);
|
||||
}
|
||||
|
||||
TEST(CUDAFunctionsTest, SynchronizeRejectsInvalidNegativeDevice) {
|
||||
if (!torch::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
ASSERT_THROW(torch::cuda::synchronize(-2), std::exception);
|
||||
}
|
||||
|
||||
TEST(CUDAFunctionsTest, CUDAGuardRestoresOriginalDeviceAfterMultipleSwitches) {
|
||||
if (!torch::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
if (torch::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
|
||||
constexpr int original_device = 0;
|
||||
constexpr int intermediate_device = 1;
|
||||
phi::backends::gpu::SetDeviceId(original_device);
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
|
||||
{
|
||||
c10::cuda::CUDAGuard guard(
|
||||
static_cast<c10::DeviceIndex>(intermediate_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), intermediate_device);
|
||||
guard.set_index(static_cast<c10::DeviceIndex>(original_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
guard.set_index(static_cast<c10::DeviceIndex>(intermediate_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), intermediate_device);
|
||||
}
|
||||
|
||||
EXPECT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
}
|
||||
|
||||
TEST(CUDAFunctionsTest,
|
||||
CUDAGuardRestoresOriginalDeviceAfterReturnToOriginalThenExit) {
|
||||
if (!torch::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
if (torch::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
|
||||
constexpr int original_device = 0;
|
||||
constexpr int intermediate_device = 1;
|
||||
phi::backends::gpu::SetDeviceId(original_device);
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
|
||||
{
|
||||
c10::cuda::CUDAGuard guard(
|
||||
static_cast<c10::DeviceIndex>(intermediate_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), intermediate_device);
|
||||
|
||||
guard.set_index(static_cast<c10::DeviceIndex>(original_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
}
|
||||
|
||||
EXPECT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
}
|
||||
|
||||
TEST(CUDAFunctionsTest,
|
||||
OptionalCUDAGuardResetRestoresOriginalDeviceAfterReturnToOriginal) {
|
||||
if (!torch::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
if (torch::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
|
||||
constexpr int original_device = 0;
|
||||
constexpr int intermediate_device = 1;
|
||||
phi::backends::gpu::SetDeviceId(original_device);
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
|
||||
c10::cuda::OptionalCUDAGuard guard;
|
||||
guard.set_index(static_cast<c10::DeviceIndex>(intermediate_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), intermediate_device);
|
||||
|
||||
guard.set_index(static_cast<c10::DeviceIndex>(original_device));
|
||||
ASSERT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
|
||||
guard.reset();
|
||||
|
||||
EXPECT_EQ(phi::backends::gpu::GetCurrentDeviceId(), original_device);
|
||||
EXPECT_FALSE(guard.original_device().has_value());
|
||||
EXPECT_FALSE(guard.current_device().has_value());
|
||||
}
|
||||
#endif
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CUDAContextLight.h — covers the 1 missing line: is_available()
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
TEST(CUDAContextLightTest, IsAvailable) {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// With GPU compilation and at least one device, this must be true.
|
||||
int gpu_count = phi::backends::gpu::GetGPUDeviceCount();
|
||||
ASSERT_EQ(at::cuda::is_available(), gpu_count > 0);
|
||||
#else
|
||||
// In CPU-only builds, is_available() should return false
|
||||
ASSERT_FALSE(at::cuda::is_available());
|
||||
#endif
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CUDAContextLight.cpp — covers all 42 missing lines
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// getNumGPUs() delegages to c10::cuda::device_count()
|
||||
TEST(CUDAContextLightTest, GetNumGPUs) {
|
||||
int64_t n = at::cuda::getNumGPUs();
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
ASSERT_EQ(n, c10::cuda::device_count());
|
||||
ASSERT_GE(n, 0);
|
||||
#else
|
||||
// In CPU-only builds, device_count() returns 0
|
||||
ASSERT_EQ(n, 0);
|
||||
#endif
|
||||
}
|
||||
|
||||
// CPU-only: device_count() must return 0 instead of throwing, matching the
|
||||
// PyTorch contract that device_count() is a non-throwing query.
|
||||
TEST(CUDAContextLightTest, DeviceCountReturnsZeroInCpuOnly) {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// Only relevant in CPU-only builds
|
||||
return;
|
||||
#else
|
||||
ASSERT_NO_THROW({
|
||||
EXPECT_EQ(c10::cuda::device_count(), 0);
|
||||
EXPECT_EQ(torch::cuda::device_count(), 0);
|
||||
});
|
||||
#endif
|
||||
}
|
||||
|
||||
// CPU-only: is_available() must be false and not throw, matching PyTorch.
|
||||
TEST(CUDAContextLightTest, IsAvailableFalseAndNoThrowInCpuOnly) {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// Only relevant in CPU-only builds
|
||||
return;
|
||||
#else
|
||||
ASSERT_NO_THROW({
|
||||
EXPECT_FALSE(at::cuda::is_available());
|
||||
EXPECT_FALSE(torch::cuda::is_available());
|
||||
});
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
|
||||
// The following tests require CUDA runtime and can only run in CUDA builds
|
||||
|
||||
// getCurrentDeviceProperties() / getDeviceProperties()
|
||||
TEST(CUDAContextLightTest, DeviceProperties) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::cuda::CUDAContextDeviceProp* prop =
|
||||
at::cuda::getCurrentDeviceProperties();
|
||||
ASSERT_NE(prop, nullptr);
|
||||
// Sanity-check a few well-known fields
|
||||
ASSERT_GT(prop->multiProcessorCount, 0);
|
||||
ASSERT_GT(prop->totalGlobalMem, 0UL);
|
||||
|
||||
// getDeviceProperties(explicit device id) must return the same struct
|
||||
int device_id = phi::backends::gpu::GetCurrentDeviceId();
|
||||
at::cuda::CUDAContextDeviceProp* prop2 =
|
||||
at::cuda::getDeviceProperties(device_id);
|
||||
ASSERT_EQ(prop, prop2);
|
||||
}
|
||||
|
||||
// warp_size()
|
||||
TEST(CUDAContextLightTest, WarpSize) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
int ws = at::cuda::warp_size();
|
||||
// All NVIDIA and AMD GPU architectures have warp size of 32 or 64
|
||||
ASSERT_TRUE(ws == 32 || ws == 64);
|
||||
}
|
||||
|
||||
// canDeviceAccessPeer() — a device cannot peer-access itself
|
||||
TEST(CUDAContextLightTest, CanDeviceAccessPeer) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
int device_id = phi::backends::gpu::GetCurrentDeviceId();
|
||||
// Self-to-self peer access is always false per CUDA spec
|
||||
bool self_peer = at::cuda::canDeviceAccessPeer(device_id, device_id);
|
||||
ASSERT_FALSE(self_peer);
|
||||
}
|
||||
|
||||
// Handle accessors — all must return non-null handles
|
||||
TEST(CUDAContextLightTest, GetCurrentCUDABlasHandle) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::cuda::CUDAContextBlasHandle h = at::cuda::getCurrentCUDABlasHandle();
|
||||
ASSERT_NE(h, nullptr);
|
||||
}
|
||||
|
||||
TEST(CUDAContextLightTest, GetCurrentCUDABlasLtHandle) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::cuda::CUDAContextBlasLtHandle h = at::cuda::getCurrentCUDABlasLtHandle();
|
||||
ASSERT_NE(h, nullptr);
|
||||
}
|
||||
|
||||
TEST(CUDAContextLightTest, GetCurrentCUDASparseHandle) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::cuda::CUDAContextSparseHandle h = at::cuda::getCurrentCUDASparseHandle();
|
||||
ASSERT_NE(h, nullptr);
|
||||
}
|
||||
|
||||
#if defined(CUDART_VERSION) || defined(USE_ROCM)
|
||||
TEST(CUDAContextLightTest, GetCurrentCUDASolverDnHandle) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::cuda::CUDAContextSolverHandle h =
|
||||
at::cuda::getCurrentCUDASolverDnHandle();
|
||||
ASSERT_NE(h, nullptr);
|
||||
}
|
||||
#endif
|
||||
|
||||
// clearCublasWorkspaces() — must not crash (no-op in the compat layer)
|
||||
TEST(CUDAContextLightTest, ClearCublasWorkspaces) {
|
||||
ASSERT_NO_THROW(at::cuda::clearCublasWorkspaces());
|
||||
}
|
||||
|
||||
// cublas_handle_stream_to_workspace() — must return a stable reference
|
||||
TEST(CUDAContextLightTest, CublasHandleStreamToWorkspace) {
|
||||
at::cuda::WorkspaceMapWithMutex& wm =
|
||||
at::cuda::cublas_handle_stream_to_workspace();
|
||||
// The map should start empty
|
||||
ASSERT_TRUE(wm.map.empty());
|
||||
// Two calls must return the same singleton
|
||||
ASSERT_EQ(&wm, &at::cuda::cublas_handle_stream_to_workspace());
|
||||
}
|
||||
|
||||
// cublaslt_handle_stream_to_workspace() — same contract
|
||||
TEST(CUDAContextLightTest, CublasLtHandleStreamToWorkspace) {
|
||||
at::cuda::WorkspaceMapWithMutex& wm =
|
||||
at::cuda::cublaslt_handle_stream_to_workspace();
|
||||
ASSERT_TRUE(wm.map.empty());
|
||||
ASSERT_EQ(&wm, &at::cuda::cublaslt_handle_stream_to_workspace());
|
||||
}
|
||||
|
||||
// getChosenWorkspaceSize() — must be 32 MiB
|
||||
TEST(CUDAContextLightTest, GetChosenWorkspaceSize) {
|
||||
constexpr size_t kExpected = 32UL * 1024UL * 1024UL;
|
||||
ASSERT_EQ(at::cuda::getChosenWorkspaceSize(), kExpected);
|
||||
}
|
||||
|
||||
// getCUDABlasLtWorkspaceSize() / getCUDABlasLtWorkspace()
|
||||
TEST(CUDAContextLightTest, CUDABlasLtWorkspace) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
size_t sz = at::cuda::getCUDABlasLtWorkspaceSize();
|
||||
ASSERT_GT(sz, 0UL);
|
||||
|
||||
void* ptr = at::cuda::getCUDABlasLtWorkspace();
|
||||
ASSERT_NE(ptr, nullptr);
|
||||
}
|
||||
|
||||
TEST(CUDAContextLightTest, CUDADeviceAllocatorSingleton) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
c10::Allocator* a0 = at::cuda::getCUDADeviceAllocator();
|
||||
c10::Allocator* a1 = at::cuda::getCUDADeviceAllocator();
|
||||
ASSERT_NE(a0, nullptr);
|
||||
ASSERT_EQ(a0, a1);
|
||||
}
|
||||
|
||||
TEST(CUDAContextLightTest, CUDADeviceAllocatorCloneAndCopyData) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
c10::Allocator* alloc = at::cuda::getCUDADeviceAllocator();
|
||||
ASSERT_NE(alloc, nullptr);
|
||||
|
||||
constexpr size_t kBytes = 32;
|
||||
c10::DataPtr src = alloc->allocate(kBytes);
|
||||
ASSERT_NE(src.get(), nullptr);
|
||||
|
||||
uint8_t h_src[kBytes];
|
||||
uint8_t h_dst[kBytes];
|
||||
for (size_t i = 0; i < kBytes; ++i) {
|
||||
h_src[i] = static_cast<uint8_t>(i + 1);
|
||||
h_dst[i] = 0;
|
||||
}
|
||||
|
||||
ASSERT_EQ(MEMCPY_FN(src.get(), h_src, kBytes, MEMCPY_HOST_TO_DEVICE),
|
||||
SUCCESS_CODE);
|
||||
|
||||
c10::DataPtr cloned = alloc->clone(src.get(), kBytes);
|
||||
ASSERT_NE(cloned.get(), nullptr);
|
||||
|
||||
ASSERT_EQ(MEMCPY_FN(h_dst, cloned.get(), kBytes, MEMCPY_DEVICE_TO_HOST),
|
||||
SUCCESS_CODE);
|
||||
ASSERT_EQ(DEVICE_SYNCHRONIZE_FN(), SUCCESS_CODE);
|
||||
|
||||
for (size_t i = 0; i < kBytes; ++i) {
|
||||
ASSERT_EQ(h_dst[i], h_src[i]);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(CUDAContextLightTest, CUDADeviceAllocatorCloneZeroBytes) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
c10::Allocator* alloc = at::cuda::getCUDADeviceAllocator();
|
||||
ASSERT_NE(alloc, nullptr);
|
||||
|
||||
c10::DataPtr src = alloc->allocate(0);
|
||||
ASSERT_EQ(src.get(), nullptr);
|
||||
|
||||
c10::DataPtr cloned = alloc->clone(src.get(), 0);
|
||||
ASSERT_EQ(cloned.get(), nullptr);
|
||||
ASSERT_EQ(cloned.device().type(), c10::DeviceType::CUDA);
|
||||
}
|
||||
|
||||
TEST(CUDAContextLightTest, AllocatorZeroSizeAndNoopCopyBranches) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
c10::Allocator* alloc = at::cuda::getCUDADeviceAllocator();
|
||||
ASSERT_NE(alloc, nullptr);
|
||||
|
||||
c10::DataPtr zero = alloc->allocate(0);
|
||||
ASSERT_EQ(zero.device().type(), c10::DeviceType::CUDA);
|
||||
ASSERT_EQ(alloc->raw_deleter(), nullptr);
|
||||
|
||||
// n==0 branch should early-return without touching pointers.
|
||||
alloc->copy_data(nullptr, nullptr, 0);
|
||||
}
|
||||
|
||||
#if defined(USE_CUDSS)
|
||||
TEST(CUDAContextLightTest, CudssHandleIsUnimplemented) {
|
||||
ASSERT_THROW((void)at::cuda::getCurrentCudssHandle(), std::exception);
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
@@ -0,0 +1,220 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// Test for TensorBase::accessor()
|
||||
TEST(TensorAccessorTest, AccessorBasic) {
|
||||
// Create a 2D tensor with known values
|
||||
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
// Get accessor
|
||||
auto accessor = tensor.accessor<float, 2>();
|
||||
|
||||
// Verify accessor dimensions
|
||||
ASSERT_EQ(accessor.size(0), 3);
|
||||
ASSERT_EQ(accessor.size(1), 4);
|
||||
|
||||
// Verify accessor values
|
||||
float expected = 0.0f;
|
||||
for (int64_t i = 0; i < 3; ++i) {
|
||||
for (int64_t j = 0; j < 4; ++j) {
|
||||
ASSERT_EQ(accessor[i][j], expected);
|
||||
expected += 1.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorAccessorTest, AccessorWithConstType) {
|
||||
// Create a tensor
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Get const accessor
|
||||
auto accessor = tensor.accessor<const float, 2>();
|
||||
|
||||
// Verify values are all ones
|
||||
for (int64_t i = 0; i < 2; ++i) {
|
||||
for (int64_t j = 0; j < 3; ++j) {
|
||||
ASSERT_EQ(accessor[i][j], 1.0f);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorAccessorTest, Accessor3D) {
|
||||
// Create a 3D tensor
|
||||
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
||||
|
||||
// Get accessor
|
||||
auto accessor = tensor.accessor<float, 3>();
|
||||
|
||||
// Verify dimensions
|
||||
ASSERT_EQ(accessor.size(0), 2);
|
||||
ASSERT_EQ(accessor.size(1), 3);
|
||||
ASSERT_EQ(accessor.size(2), 4);
|
||||
|
||||
// Verify a few values
|
||||
ASSERT_EQ(accessor[0][0][0], 0.0f);
|
||||
ASSERT_EQ(accessor[0][0][3], 3.0f);
|
||||
ASSERT_EQ(accessor[1][2][3], 23.0f);
|
||||
}
|
||||
|
||||
TEST(TensorAccessorTest, AccessorModifyValues) {
|
||||
// Create a tensor
|
||||
at::Tensor tensor = at::zeros({2, 3}, at::kFloat);
|
||||
|
||||
// Get mutable accessor
|
||||
auto accessor = tensor.accessor<float, 2>();
|
||||
|
||||
// Modify values through accessor
|
||||
for (int64_t i = 0; i < 2; ++i) {
|
||||
for (int64_t j = 0; j < 3; ++j) {
|
||||
accessor[i][j] = static_cast<float>(i * 3 + j);
|
||||
}
|
||||
}
|
||||
|
||||
// Verify modifications via data_ptr
|
||||
float* data = tensor.data_ptr<float>();
|
||||
for (int64_t i = 0; i < 6; ++i) {
|
||||
ASSERT_EQ(data[i], static_cast<float>(i));
|
||||
}
|
||||
}
|
||||
|
||||
// Test for TensorBase::packed_accessor64()
|
||||
TEST(TensorAccessorTest, PackedAccessor64Basic) {
|
||||
// Create a 2D tensor
|
||||
at::Tensor tensor = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
// Get packed accessor with int64_t index type
|
||||
auto packed = tensor.packed_accessor64<float, 2>();
|
||||
|
||||
// Verify dimensions
|
||||
ASSERT_EQ(packed.size(0), 3);
|
||||
ASSERT_EQ(packed.size(1), 4);
|
||||
|
||||
// Verify strides
|
||||
ASSERT_EQ(packed.stride(0), 4);
|
||||
ASSERT_EQ(packed.stride(1), 1);
|
||||
|
||||
// Verify values
|
||||
float expected = 0.0f;
|
||||
for (int64_t i = 0; i < 3; ++i) {
|
||||
for (int64_t j = 0; j < 4; ++j) {
|
||||
ASSERT_EQ(packed[i][j], expected);
|
||||
expected += 1.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Test for TensorBase::packed_accessor32()
|
||||
TEST(TensorAccessorTest, PackedAccessor32Basic) {
|
||||
// Create a small 2D tensor (within int32_t range)
|
||||
at::Tensor tensor = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
// Get packed accessor with int32_t index type
|
||||
auto packed = tensor.packed_accessor32<float, 2>();
|
||||
|
||||
// Verify dimensions
|
||||
ASSERT_EQ(packed.size(0), 2);
|
||||
ASSERT_EQ(packed.size(1), 3);
|
||||
|
||||
// Verify strides
|
||||
ASSERT_EQ(packed.stride(0), 3);
|
||||
ASSERT_EQ(packed.stride(1), 1);
|
||||
|
||||
// Verify values
|
||||
ASSERT_EQ(packed[0][0], 0.0f);
|
||||
ASSERT_EQ(packed[0][2], 2.0f);
|
||||
ASSERT_EQ(packed[1][0], 3.0f);
|
||||
ASSERT_EQ(packed[1][2], 5.0f);
|
||||
}
|
||||
|
||||
// Test for TensorBase::generic_packed_accessor()
|
||||
TEST(TensorAccessorTest, GenericPackedAccessor) {
|
||||
// Create a 3D tensor
|
||||
at::Tensor tensor = at::arange(24, at::kDouble).reshape({2, 3, 4});
|
||||
|
||||
// Get generic packed accessor with default template parameters
|
||||
auto packed = tensor.generic_packed_accessor<double, 3>();
|
||||
|
||||
// Verify dimensions
|
||||
ASSERT_EQ(packed.size(0), 2);
|
||||
ASSERT_EQ(packed.size(1), 3);
|
||||
ASSERT_EQ(packed.size(2), 4);
|
||||
|
||||
// Verify strides
|
||||
ASSERT_EQ(packed.stride(0), 12); // 3*4
|
||||
ASSERT_EQ(packed.stride(1), 4);
|
||||
ASSERT_EQ(packed.stride(2), 1);
|
||||
|
||||
// Verify corner values
|
||||
ASSERT_DOUBLE_EQ(packed[0][0][0], 0.0);
|
||||
ASSERT_DOUBLE_EQ(packed[1][2][3], 23.0);
|
||||
}
|
||||
|
||||
TEST(TensorAccessorTest, PackedAccessorWithIntType) {
|
||||
// Test with integer tensor
|
||||
at::Tensor tensor = at::arange(10, at::kInt).reshape({2, 5});
|
||||
|
||||
auto packed = tensor.packed_accessor64<int, 2>();
|
||||
|
||||
ASSERT_EQ(packed.size(0), 2);
|
||||
ASSERT_EQ(packed.size(1), 5);
|
||||
|
||||
int expected = 0;
|
||||
for (int64_t i = 0; i < 2; ++i) {
|
||||
for (int64_t j = 0; j < 5; ++j) {
|
||||
ASSERT_EQ(packed[i][j], expected);
|
||||
expected++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(TensorAccessorTest, PackedAccessorCUDA) {
|
||||
if (at::cuda::is_available()) {
|
||||
// Create CUDA tensor
|
||||
at::Tensor tensor =
|
||||
at::arange(12, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA))
|
||||
.reshape({3, 4});
|
||||
|
||||
// Get packed accessor (typically used to pass to CUDA kernels)
|
||||
auto packed = tensor.packed_accessor64<float, 2>();
|
||||
|
||||
// Verify dimensions
|
||||
ASSERT_EQ(packed.size(0), 3);
|
||||
ASSERT_EQ(packed.size(1), 4);
|
||||
|
||||
// Verify strides
|
||||
ASSERT_EQ(packed.stride(0), 4);
|
||||
ASSERT_EQ(packed.stride(1), 1);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,181 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/Utils.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::detail::tensor_cpu / tensor_backend / complex variants
|
||||
// and the at::tensor() factory macro-generated overloads (ATen/Utils.h)
|
||||
// ============================================================
|
||||
|
||||
// ---- tensor_cpu (via at::tensor public API) ----
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Float) {
|
||||
std::vector<float> data = {1.0f, 2.0f, 3.0f};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<float>(data),
|
||||
at::TensorOptions().dtype(at::kFloat));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_EQ(t.numel(), 3);
|
||||
ASSERT_NEAR(t[0].item<float>(), 1.0f, 1e-6f);
|
||||
ASSERT_NEAR(t[2].item<float>(), 3.0f, 1e-6f);
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Double) {
|
||||
std::vector<double> data = {1.1, 2.2, 3.3};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<double>(data),
|
||||
at::TensorOptions().dtype(at::kDouble));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kDouble);
|
||||
ASSERT_NEAR(t[1].item<double>(), 2.2, 1e-10);
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Int32) {
|
||||
std::vector<int32_t> data = {10, 20, 30};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<int32_t>(data),
|
||||
at::TensorOptions().dtype(at::kInt));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kInt);
|
||||
ASSERT_EQ(t[0].item<int32_t>(), 10);
|
||||
ASSERT_EQ(t[2].item<int32_t>(), 30);
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Int64) {
|
||||
std::vector<int64_t> data = {100LL, 200LL};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<int64_t>(data),
|
||||
at::TensorOptions().dtype(at::kLong));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(t[1].item<int64_t>(), 200LL);
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Int8) {
|
||||
std::vector<int8_t> data = {-1, 0, 1};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<int8_t>(data),
|
||||
at::TensorOptions().dtype(at::kChar));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kChar);
|
||||
ASSERT_EQ(t[0].item<int8_t>(), static_cast<int8_t>(-1));
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Int16) {
|
||||
std::vector<int16_t> data = {256, 512};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<int16_t>(data),
|
||||
at::TensorOptions().dtype(at::kShort));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kShort);
|
||||
ASSERT_EQ(t[0].item<int16_t>(), static_cast<int16_t>(256));
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_UInt8) {
|
||||
std::vector<uint8_t> data = {200, 255};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<uint8_t>(data),
|
||||
at::TensorOptions().dtype(at::kByte));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kByte);
|
||||
ASSERT_EQ(t[1].item<uint8_t>(), static_cast<uint8_t>(255));
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_Bool) {
|
||||
// std::vector<bool> is a bitfield specialization without data(), so use a
|
||||
// plain C array to construct c10::ArrayRef<bool>.
|
||||
bool data[] = {true, false, true};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<bool>(data),
|
||||
at::TensorOptions().dtype(at::kBool));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kBool);
|
||||
ASSERT_TRUE(t[0].item<bool>());
|
||||
ASSERT_FALSE(t[1].item<bool>());
|
||||
}
|
||||
|
||||
// ---- dtype promotion: values stored as native type then cast ----
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_DtypePromotion_IntToFloat) {
|
||||
// Store int32 values, but request float32 output – should auto-cast.
|
||||
std::vector<int32_t> data = {1, 2, 3};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<int32_t>(data),
|
||||
at::TensorOptions().dtype(at::kFloat));
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_NEAR(t[0].item<float>(), 1.0f, 1e-6f);
|
||||
}
|
||||
|
||||
// ---- contiguity ----
|
||||
|
||||
TEST(ATenUtilsTest, TensorCPU_IsContiguous) {
|
||||
std::vector<float> data = {1.0f, 2.0f, 3.0f, 4.0f};
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<float>(data),
|
||||
at::TensorOptions().dtype(at::kFloat));
|
||||
|
||||
ASSERT_TRUE(t.is_contiguous());
|
||||
}
|
||||
|
||||
// ---- tensor_backend (CPU -> same result since default is CPU in tests) ----
|
||||
|
||||
TEST(ATenUtilsTest, TensorBackend_CPUDevice_MatchesTensorCPU) {
|
||||
std::vector<float> data = {5.0f, 6.0f};
|
||||
at::TensorOptions opts =
|
||||
at::TensorOptions().dtype(at::kFloat).device(c10::Device(c10::kCPU));
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<float>(data), opts);
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_EQ(t.device().type(), c10::DeviceType::CPU);
|
||||
ASSERT_NEAR(t[0].item<float>(), 5.0f, 1e-6f);
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(ATenUtilsTest, TensorBackend_GPUDevice) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
std::vector<float> data = {7.0f, 8.0f};
|
||||
at::TensorOptions opts =
|
||||
at::TensorOptions().dtype(at::kFloat).device(c10::Device(c10::kCUDA, 0));
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<float>(data), opts);
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_EQ(t.device().type(), c10::DeviceType::CUDA);
|
||||
}
|
||||
|
||||
TEST(ATenUtilsTest, TensorComplexBackend_GPUDevice) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
std::vector<c10::complex<float>> data = {{1.0f, 0.0f}};
|
||||
at::TensorOptions opts = at::TensorOptions()
|
||||
.dtype(at::kComplexFloat)
|
||||
.device(c10::Device(c10::kCUDA, 0));
|
||||
at::Tensor t = at::tensor(c10::ArrayRef<c10::complex<float>>(data), opts);
|
||||
|
||||
ASSERT_EQ(t.scalar_type(), at::kComplexFloat);
|
||||
ASSERT_EQ(t.device().type(), c10::DeviceType::CUDA);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,656 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <limits>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
TEST(TestAll, AllNoDim) {
|
||||
// Test all() without arguments - check all elements in tensor
|
||||
at::Tensor tensor = at::ones({3}, at::kBool);
|
||||
tensor[1] = false;
|
||||
at::Tensor result = tensor.all();
|
||||
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_EQ(result.item<bool>(), false);
|
||||
|
||||
// Test with all true values
|
||||
at::Tensor tensor_all_true = at::ones({3}, at::kBool);
|
||||
at::Tensor result_all_true = tensor_all_true.all();
|
||||
ASSERT_EQ(result_all_true.item<bool>(), true);
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithDim) {
|
||||
// Test all(dim) - check along specific dimension
|
||||
at::Tensor tensor = at::ones({2, 2}, at::kBool);
|
||||
tensor[1][0] = false;
|
||||
|
||||
// All along dimension 0
|
||||
at::Tensor result_dim0 = tensor.all(0);
|
||||
ASSERT_EQ(result_dim0.sizes(), c10::IntArrayRef({2}));
|
||||
ASSERT_EQ(result_dim0.data_ptr<bool>()[0], false); // column 0 has false
|
||||
ASSERT_EQ(result_dim0.data_ptr<bool>()[1], true); // column 1 has all true
|
||||
|
||||
// All along dimension 1
|
||||
at::Tensor result_dim1 = tensor.all(1);
|
||||
ASSERT_EQ(result_dim1.sizes(), c10::IntArrayRef({2}));
|
||||
ASSERT_EQ(result_dim1.data_ptr<bool>()[0], true); // row 0 has all true
|
||||
ASSERT_EQ(result_dim1.data_ptr<bool>()[1], false); // row 1 has false
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithDimKeepdim) {
|
||||
// Test all(dim, keepdim) - keep the dimension
|
||||
at::Tensor tensor = at::ones({2, 2}, at::kBool);
|
||||
|
||||
at::Tensor result = tensor.all(0, true);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({1, 2}));
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithOptionalDim) {
|
||||
// Test all(OptionalIntArrayRef dim, keepdim)
|
||||
at::Tensor tensor = at::ones({2, 2}, at::kBool);
|
||||
|
||||
// With specific dimensions
|
||||
at::Tensor result = tensor.all(c10::IntArrayRef({0}), false);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2}));
|
||||
}
|
||||
|
||||
TEST(TestAll, AllNoDimAllFalse) {
|
||||
// Test all() on tensor with all false values
|
||||
at::Tensor tensor = at::zeros({4}, at::kBool);
|
||||
at::Tensor result = tensor.all();
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_EQ(result.item<bool>(), false);
|
||||
}
|
||||
|
||||
TEST(TestAll, AllNoDimSingleElement) {
|
||||
// Test all() on single-element tensor
|
||||
at::Tensor tensor_true = at::ones({1}, at::kBool);
|
||||
ASSERT_EQ(tensor_true.all().item<bool>(), true);
|
||||
|
||||
at::Tensor tensor_false = at::zeros({1}, at::kBool);
|
||||
ASSERT_EQ(tensor_false.all().item<bool>(), false);
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithNegativeDim) {
|
||||
// Test all(dim) with negative dimension index
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
tensor[0][1] = false;
|
||||
at::Tensor result = tensor.all(-1); // equivalent to dim=1
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2}));
|
||||
ASSERT_EQ(result.data_ptr<bool>()[0], false); // row 0 has a false
|
||||
ASSERT_EQ(result.data_ptr<bool>()[1], true); // row 1 all true
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithDimKeepdimTrue) {
|
||||
// Test all(dim, keepdim=true) with different dims
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
tensor[1][0] = false;
|
||||
|
||||
at::Tensor result_dim0 = tensor.all(0, true);
|
||||
ASSERT_EQ(result_dim0.sizes(), c10::IntArrayRef({1, 3}));
|
||||
ASSERT_EQ(result_dim0.data_ptr<bool>()[0], false); // col 0 has false
|
||||
ASSERT_EQ(result_dim0.data_ptr<bool>()[1], true);
|
||||
ASSERT_EQ(result_dim0.data_ptr<bool>()[2], true);
|
||||
|
||||
at::Tensor result_dim1 = tensor.all(1, true);
|
||||
ASSERT_EQ(result_dim1.sizes(), c10::IntArrayRef({2, 1}));
|
||||
ASSERT_EQ(result_dim1.data_ptr<bool>()[0], true); // row 0 all true
|
||||
ASSERT_EQ(result_dim1.data_ptr<bool>()[1], false); // row 1 has false
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithOptionalDimNullopt) {
|
||||
// Test all(OptionalIntArrayRef) with nullopt - reduces all dimensions
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
at::OptionalIntArrayRef dim = std::nullopt;
|
||||
at::Tensor result = at::all(tensor, dim, false);
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_EQ(result.item<bool>(), true);
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithOptionalDimNulloptHasFalse) {
|
||||
// Test all(OptionalIntArrayRef nullopt) when tensor contains false
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
tensor[1][2] = false;
|
||||
at::OptionalIntArrayRef dim = std::nullopt;
|
||||
at::Tensor result = at::all(tensor, dim, false);
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_EQ(result.item<bool>(), false);
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithOptionalDimKeepdim) {
|
||||
// Test all(OptionalIntArrayRef, keepdim=true)
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
at::Tensor result = at::all(tensor, c10::IntArrayRef({0}), true);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({1, 3}));
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWithOptionalMultipleDims) {
|
||||
// Test all(OptionalIntArrayRef) with multiple dimensions
|
||||
at::Tensor tensor = at::ones({2, 3, 4}, at::kBool);
|
||||
at::Tensor result = at::all(tensor, c10::IntArrayRef({0, 2}), false);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({3}));
|
||||
// All elements are true, so result should be all true
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
ASSERT_EQ(result.data_ptr<bool>()[i], true);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestAll, MemberAllWithOptionalNullopt) {
|
||||
// Test member function Tensor::all(OptionalIntArrayRef, keepdim) with nullopt
|
||||
at::Tensor tensor = at::ones({3, 4}, at::kBool);
|
||||
at::OptionalIntArrayRef dim = std::nullopt;
|
||||
at::Tensor result = tensor.all(dim, false);
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_EQ(result.item<bool>(), true);
|
||||
}
|
||||
|
||||
TEST(TestAll, MemberAllWithOptionalNulloptKeepdim) {
|
||||
// Test member function Tensor::all(nullopt, keepdim=true)
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
at::OptionalIntArrayRef dim = std::nullopt;
|
||||
at::Tensor result = tensor.all(dim, true);
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_EQ(result.item<bool>(), true);
|
||||
}
|
||||
|
||||
TEST(TestAll, StandaloneFunction) {
|
||||
// Test at::all() standalone function
|
||||
at::Tensor tensor = at::ones({3}, at::kBool);
|
||||
tensor[2] = false;
|
||||
at::Tensor result = at::all(tensor);
|
||||
|
||||
ASSERT_EQ(result.item<bool>(), false);
|
||||
}
|
||||
|
||||
TEST(TestAll, StandaloneFunctionWithDim) {
|
||||
// Test at::all(tensor, dim, keepdim)
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kBool);
|
||||
tensor[0][0] = false;
|
||||
|
||||
at::Tensor result = at::all(tensor, 0, false);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({3}));
|
||||
ASSERT_EQ(result.data_ptr<bool>()[0], false);
|
||||
ASSERT_EQ(result.data_ptr<bool>()[1], true);
|
||||
ASSERT_EQ(result.data_ptr<bool>()[2], true);
|
||||
|
||||
at::Tensor result_kd = at::all(tensor, 0, true);
|
||||
ASSERT_EQ(result_kd.sizes(), c10::IntArrayRef({1, 3}));
|
||||
}
|
||||
|
||||
TEST(TestAll, AllWith3DTensor) {
|
||||
// Test all on a 3D tensor to exercise more paths
|
||||
at::Tensor tensor = at::ones({2, 2, 2}, at::kBool);
|
||||
tensor[0][0][0] = false;
|
||||
|
||||
at::Tensor result_all = tensor.all();
|
||||
ASSERT_EQ(result_all.item<bool>(), false);
|
||||
|
||||
at::Tensor result_dim0 = tensor.all(0, false);
|
||||
ASSERT_EQ(result_dim0.sizes(), c10::IntArrayRef({2, 2}));
|
||||
|
||||
at::Tensor result_dim2 = tensor.all(2, true);
|
||||
ASSERT_EQ(result_dim2.sizes(), c10::IntArrayRef({2, 2, 1}));
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseBasic) {
|
||||
// Test allclose - basic equal tensors
|
||||
at::Tensor tensor1 = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor tensor2 = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
bool result = tensor1.allclose(tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseNotEqual) {
|
||||
// Test allclose - tensors that are not close
|
||||
at::Tensor tensor1 = at::arange(1, 4, at::TensorOptions().dtype(at::kFloat));
|
||||
at::Tensor tensor2 = tensor1.clone();
|
||||
tensor2[2] = 4.0f;
|
||||
|
||||
bool result = tensor1.allclose(tensor2);
|
||||
ASSERT_EQ(result, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, StandaloneFunction) {
|
||||
// Test at::allclose() standalone function
|
||||
at::Tensor tensor1 = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor tensor2 = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseWithCustomRtol) {
|
||||
// Test allclose with custom relative tolerance
|
||||
at::Tensor tensor1 = at::ones({3}, at::kFloat);
|
||||
at::Tensor tensor2 = at::ones({3}, at::kFloat);
|
||||
tensor2[0] = 1.05f; // 5% difference
|
||||
|
||||
// With default rtol=1e-05, should fail
|
||||
bool result_default = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result_default, false);
|
||||
|
||||
// With rtol=0.1 (10%), 5% difference should pass
|
||||
bool result_large_rtol = at::allclose(tensor1, tensor2, 0.1, 1e-08, false);
|
||||
ASSERT_EQ(result_large_rtol, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseWithCustomAtol) {
|
||||
// Test allclose with custom absolute tolerance
|
||||
at::Tensor tensor1 = at::zeros({3}, at::kFloat);
|
||||
at::Tensor tensor2 = at::zeros({3}, at::kFloat);
|
||||
tensor2[1] = 0.05f;
|
||||
|
||||
// With default atol=1e-08, should fail
|
||||
bool result_default = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result_default, false);
|
||||
|
||||
// With atol=0.1, should pass
|
||||
bool result_large_atol = at::allclose(tensor1, tensor2, 1e-05, 0.1, false);
|
||||
ASSERT_EQ(result_large_atol, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseMemberWithAllParams) {
|
||||
// Test Tensor::allclose member function with all explicit parameters
|
||||
at::Tensor tensor1 = at::ones({2, 2}, at::kFloat);
|
||||
at::Tensor tensor2 = at::ones({2, 2}, at::kFloat);
|
||||
|
||||
bool result = tensor1.allclose(tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseMemberNotClose) {
|
||||
// Test Tensor::allclose member function returns false when not close
|
||||
at::Tensor tensor1 = at::ones({2, 3}, at::kFloat);
|
||||
at::Tensor tensor2 = at::ones({2, 3}, at::kFloat);
|
||||
tensor2[0][0] = 100.0f;
|
||||
|
||||
bool result = tensor1.allclose(tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseMemberWithCustomTolerance) {
|
||||
// Test Tensor::allclose member function with custom rtol and atol
|
||||
at::Tensor tensor1 = at::ones({4}, at::kFloat);
|
||||
at::Tensor tensor2 = at::ones({4}, at::kFloat);
|
||||
tensor2[3] = 1.001f; // small relative difference
|
||||
|
||||
// Default tolerance should fail
|
||||
ASSERT_EQ(tensor1.allclose(tensor2), false);
|
||||
|
||||
// Custom rtol=0.01 (1%) should pass
|
||||
ASSERT_EQ(tensor1.allclose(tensor2, 0.01, 1e-08, false), true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseExactZeros) {
|
||||
// Test allclose with exact zero tensors
|
||||
at::Tensor tensor1 = at::zeros({5}, at::kFloat);
|
||||
at::Tensor tensor2 = at::zeros({5}, at::kFloat);
|
||||
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
bool result_member = tensor1.allclose(tensor2);
|
||||
ASSERT_EQ(result_member, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseHighDim) {
|
||||
// Test allclose with higher dimensional tensors
|
||||
at::Tensor tensor1 = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
||||
at::Tensor tensor2 = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
||||
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
bool result_member = tensor1.allclose(tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_member, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseEqualNanDefaultFalse) {
|
||||
// Test allclose default behavior: NaN != NaN when equal_nan not set
|
||||
// Use from_blob to avoid triggering fill_ operation which doesn't support NaN
|
||||
const float nan_val = std::numeric_limits<float>::quiet_NaN();
|
||||
float data1[3] = {0.0f, nan_val, 0.0f};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {3}, at::kFloat);
|
||||
at::Tensor tensor2 = tensor1.clone();
|
||||
|
||||
// Default equal_nan=false: NaN is not equal to NaN, so result is false
|
||||
bool result_standalone = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result_standalone, false);
|
||||
|
||||
bool result_member = tensor1.allclose(tensor2);
|
||||
ASSERT_EQ(result_member, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseEqualNanTrue) {
|
||||
// Test allclose with equal_nan=true: NaN == NaN should yield true
|
||||
// Use from_blob to avoid triggering fill_ operation which doesn't support NaN
|
||||
const float nan_val = std::numeric_limits<float>::quiet_NaN();
|
||||
float data1[3] = {0.0f, nan_val, 0.0f};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {3}, at::kFloat);
|
||||
at::Tensor tensor2 = tensor1.clone();
|
||||
|
||||
// equal_nan=true: NaN is treated as equal to NaN
|
||||
bool result = at::allclose(tensor1, tensor2, 1e-05, 1e-08, true);
|
||||
ASSERT_EQ(result, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseEqualNanTrueAllNan) {
|
||||
// Test allclose with equal_nan=true on all-NaN tensors
|
||||
// Use from_blob to avoid triggering fill_ operation which doesn't support NaN
|
||||
const float nan_val = std::numeric_limits<float>::quiet_NaN();
|
||||
float data1[4] = {nan_val, nan_val, nan_val, nan_val};
|
||||
float data2[4] = {nan_val, nan_val, nan_val, nan_val};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {4}, at::kFloat);
|
||||
at::Tensor tensor2 = at::from_blob(data2, {4}, at::kFloat);
|
||||
|
||||
bool result_equal_nan = at::allclose(tensor1, tensor2, 1e-05, 1e-08, true);
|
||||
ASSERT_EQ(result_equal_nan, true);
|
||||
|
||||
// Without equal_nan, all-NaN tensors should not be close
|
||||
bool result_no_equal_nan =
|
||||
at::allclose(tensor1, tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_no_equal_nan, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseMemberEqualNanTrue) {
|
||||
// Test Tensor::allclose member function with equal_nan=true
|
||||
// Use from_blob to avoid triggering fill_ operation which doesn't support NaN
|
||||
const float nan_val = std::numeric_limits<float>::quiet_NaN();
|
||||
float data1[4] = {nan_val, 0.0f, 0.0f, nan_val};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {4}, at::kFloat);
|
||||
at::Tensor tensor2 = tensor1.clone();
|
||||
|
||||
bool result_true = tensor1.allclose(tensor2, 1e-05, 1e-08, true);
|
||||
ASSERT_EQ(result_true, true);
|
||||
|
||||
bool result_false = tensor1.allclose(tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_false, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseMixedNanAndValues) {
|
||||
// Test allclose where some elements match and one is NaN
|
||||
// Use from_blob to avoid triggering fill_ operation which doesn't support NaN
|
||||
const float nan_val = std::numeric_limits<float>::quiet_NaN();
|
||||
float data1[4] = {1.0f, 1.0f, nan_val, 1.0f};
|
||||
float data2[4] = {1.0f, 1.0f, nan_val, 1.0f};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {4}, at::kFloat);
|
||||
at::Tensor tensor2 = at::from_blob(data2, {4}, at::kFloat);
|
||||
|
||||
// NaN-aware comparison: non-NaN elements are equal, NaN treated equal
|
||||
bool result_eq_nan = at::allclose(tensor1, tensor2, 1e-05, 1e-08, true);
|
||||
ASSERT_EQ(result_eq_nan, true);
|
||||
|
||||
// Without equal_nan: NaN elements fail the check
|
||||
bool result_no_eq_nan = at::allclose(tensor1, tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_no_eq_nan, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseDouble) {
|
||||
// Test allclose with double-precision (float64) tensors
|
||||
at::Tensor tensor1 = at::arange(6, at::kDouble).reshape({2, 3});
|
||||
at::Tensor tensor2 = at::arange(6, at::kDouble).reshape({2, 3});
|
||||
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
bool result_member = tensor1.allclose(tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_member, true);
|
||||
|
||||
// Introduce a small difference
|
||||
tensor2[1][2] = 5.001;
|
||||
bool result_diff = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result_diff, false);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseDoubleEqualNan) {
|
||||
// Test allclose with double-precision tensors and NaN
|
||||
// Use from_blob to avoid triggering fill_ operation which doesn't support NaN
|
||||
const double nan_val = std::numeric_limits<double>::quiet_NaN();
|
||||
double data1[3] = {nan_val, 0.0, 0.0};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {3}, at::kDouble);
|
||||
at::Tensor tensor2 = tensor1.clone();
|
||||
|
||||
bool result_false = at::allclose(tensor1, tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_false, false);
|
||||
|
||||
bool result_true = at::allclose(tensor1, tensor2, 1e-05, 1e-08, true);
|
||||
ASSERT_EQ(result_true, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseStandaloneWithExplicitParams) {
|
||||
// Test at::allclose() standalone with all explicit parameters
|
||||
at::Tensor tensor1 = at::ones({3}, at::kFloat);
|
||||
at::Tensor tensor2 = at::ones({3}, at::kFloat);
|
||||
|
||||
// All explicit parameters including equal_nan
|
||||
bool result_false_nan = at::allclose(tensor1, tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_false_nan, true);
|
||||
|
||||
bool result_true_nan = at::allclose(tensor1, tensor2, 1e-05, 1e-08, true);
|
||||
ASSERT_EQ(result_true_nan, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseInfinityValues) {
|
||||
// Test allclose with infinity values
|
||||
// Use from_blob to avoid triggering fill_ operation
|
||||
const float inf_val = std::numeric_limits<float>::infinity();
|
||||
float data1[3] = {inf_val, 1.0f, 1.0f};
|
||||
at::Tensor tensor1 = at::from_blob(data1, {3}, at::kFloat);
|
||||
at::Tensor tensor2 = tensor1.clone();
|
||||
|
||||
// Identical infinity values should be close
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
bool result_member = tensor1.allclose(tensor2, 1e-05, 1e-08, false);
|
||||
ASSERT_EQ(result_member, true);
|
||||
|
||||
// Note: PyTorch's allclose considers +inf and -inf as close because:
|
||||
// |inf - (-inf)| = inf <= (atol + rtol * |inf|) = inf
|
||||
// So this test case expectation was wrong - we just verify the behavior
|
||||
float data3[3] = {-inf_val, 1.0f, 1.0f};
|
||||
at::Tensor tensor3 = at::from_blob(data3, {3}, at::kFloat);
|
||||
bool result_diff_inf = at::allclose(tensor1, tensor3);
|
||||
// PyTorch returns true here because inf <= inf is true mathematically
|
||||
ASSERT_EQ(result_diff_inf, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseInt32) {
|
||||
// Test allclose with int32 tensors
|
||||
at::Tensor tensor1 = at::arange(6, at::kInt).reshape({2, 3});
|
||||
at::Tensor tensor2 = at::arange(6, at::kInt).reshape({2, 3});
|
||||
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
// Test with different values
|
||||
at::Tensor tensor3 = at::ones({3}, at::kInt);
|
||||
at::Tensor tensor4 = at::ones({3}, at::kInt);
|
||||
tensor4[0] = 2;
|
||||
bool result_diff = at::allclose(tensor3, tensor4);
|
||||
ASSERT_EQ(result_diff, false);
|
||||
|
||||
// Test with custom tolerance
|
||||
bool result_tol = at::allclose(tensor3, tensor4, 1.0, 0.0, false);
|
||||
ASSERT_EQ(result_tol, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseInt64) {
|
||||
// Test allclose with int64 (long) tensors
|
||||
at::Tensor tensor1 = at::arange(6, at::kLong).reshape({2, 3});
|
||||
at::Tensor tensor2 = at::arange(6, at::kLong).reshape({2, 3});
|
||||
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
// Test with small difference and custom tolerance
|
||||
at::Tensor tensor3 = at::ones({4}, at::kLong);
|
||||
at::Tensor tensor4 = at::ones({4}, at::kLong);
|
||||
tensor4[0] = 2;
|
||||
bool result_diff = at::allclose(tensor3, tensor4);
|
||||
ASSERT_EQ(result_diff, false);
|
||||
|
||||
// With large tolerance, should pass
|
||||
bool result_tol = at::allclose(tensor3, tensor4, 1.0, 0.0, false);
|
||||
ASSERT_EQ(result_tol, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseEmptyTensor) {
|
||||
// Test allclose with empty tensors
|
||||
at::Tensor tensor1 = at::empty({0}, at::kFloat);
|
||||
at::Tensor tensor2 = at::empty({0}, at::kFloat);
|
||||
|
||||
// Empty tensors should be close to each other
|
||||
bool result = at::allclose(tensor1, tensor2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
// Member function
|
||||
bool result_member = tensor1.allclose(tensor2);
|
||||
ASSERT_EQ(result_member, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseScalarTensor) {
|
||||
// Test allclose with scalar tensors (0-dimensional)
|
||||
at::Tensor scalar1 = at::tensor(1.0, at::kFloat);
|
||||
at::Tensor scalar2 = at::tensor(1.0, at::kFloat);
|
||||
|
||||
bool result = at::allclose(scalar1, scalar2);
|
||||
ASSERT_EQ(result, true);
|
||||
|
||||
// Different values
|
||||
at::Tensor scalar3 = at::tensor(1.0, at::kFloat);
|
||||
at::Tensor scalar4 = at::tensor(2.0, at::kFloat);
|
||||
bool result_diff = at::allclose(scalar3, scalar4);
|
||||
ASSERT_EQ(result_diff, false);
|
||||
|
||||
// Within tolerance
|
||||
bool result_tol = at::allclose(scalar3, scalar4, 1.0, 0.0, false);
|
||||
ASSERT_EQ(result_tol, true);
|
||||
}
|
||||
|
||||
TEST(TestAllclose, AllcloseWithDifferentRtolAtolOrder) {
|
||||
// Test allclose with parameters in different orders (edge cases)
|
||||
at::Tensor tensor1 = at::zeros({3}, at::kFloat);
|
||||
at::Tensor tensor2 = at::zeros({3}, at::kFloat);
|
||||
tensor2[0] = 0.0001f;
|
||||
|
||||
// Test with zero rtol, small atol
|
||||
bool result1 = at::allclose(tensor1, tensor2, 0.0, 0.0001, false);
|
||||
ASSERT_EQ(result1, true);
|
||||
|
||||
// Test with zero atol, small rtol
|
||||
bool result2 = at::allclose(tensor1, tensor2, 0.0001, 0.0, false);
|
||||
ASSERT_EQ(result2, false); // relative tolerance is relative to values (0.0)
|
||||
|
||||
// Both zero tolerance - exact match required
|
||||
at::Tensor tensor3 = at::ones({2}, at::kFloat);
|
||||
at::Tensor tensor4 = at::ones({2}, at::kFloat);
|
||||
bool result3 = at::allclose(tensor3, tensor4, 0.0, 0.0, false);
|
||||
ASSERT_EQ(result3, true);
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteBasic) {
|
||||
// Test absolute() - alias for abs()
|
||||
at::Tensor tensor = at::tensor({-3.0f, 2.0f, -1.0f});
|
||||
at::Tensor result = tensor.absolute();
|
||||
|
||||
ASSERT_EQ(result.numel(), 3);
|
||||
ASSERT_NEAR(result.data_ptr<float>()[0], 3.0f, 1e-6f);
|
||||
ASSERT_NEAR(result.data_ptr<float>()[1], 2.0f, 1e-6f);
|
||||
ASSERT_NEAR(result.data_ptr<float>()[2], 1.0f, 1e-6f);
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteNegativeOnly) {
|
||||
// Test absolute() on all-negative tensor
|
||||
at::Tensor tensor = at::tensor({-5.0f, -10.0f, -0.5f});
|
||||
at::Tensor result = tensor.absolute();
|
||||
|
||||
ASSERT_NEAR(result.data_ptr<float>()[0], 5.0f, 1e-6f);
|
||||
ASSERT_NEAR(result.data_ptr<float>()[1], 10.0f, 1e-6f);
|
||||
ASSERT_NEAR(result.data_ptr<float>()[2], 0.5f, 1e-6f);
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteZero) {
|
||||
// Test absolute() on zero tensor
|
||||
at::Tensor tensor = at::zeros({3}, at::kFloat);
|
||||
at::Tensor result = tensor.absolute();
|
||||
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
ASSERT_NEAR(result.data_ptr<float>()[i], 0.0f, 1e-6f);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteInPlace) {
|
||||
// Test absolute_() - in-place alias for abs_()
|
||||
at::Tensor tensor = at::tensor({-3.0f, 2.0f, -1.0f});
|
||||
at::Tensor& ref = tensor.absolute_();
|
||||
|
||||
// Should modify tensor in place
|
||||
ASSERT_NEAR(tensor.data_ptr<float>()[0], 3.0f, 1e-6f);
|
||||
ASSERT_NEAR(tensor.data_ptr<float>()[1], 2.0f, 1e-6f);
|
||||
ASSERT_NEAR(tensor.data_ptr<float>()[2], 1.0f, 1e-6f);
|
||||
|
||||
// Return value should be the same tensor
|
||||
ASSERT_EQ(ref.data_ptr<float>(), tensor.data_ptr<float>());
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteInPlaceNegative) {
|
||||
// Test absolute_() on all-negative tensor
|
||||
at::Tensor tensor = at::tensor({-4.0f, -8.0f, -0.25f});
|
||||
tensor.absolute_();
|
||||
|
||||
ASSERT_NEAR(tensor.data_ptr<float>()[0], 4.0f, 1e-6f);
|
||||
ASSERT_NEAR(tensor.data_ptr<float>()[1], 8.0f, 1e-6f);
|
||||
ASSERT_NEAR(tensor.data_ptr<float>()[2], 0.25f, 1e-6f);
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteDouble) {
|
||||
// Test absolute() with double precision
|
||||
at::Tensor tensor = at::tensor({-1.5, 2.5, -3.5}, at::kDouble);
|
||||
at::Tensor result = tensor.absolute();
|
||||
|
||||
ASSERT_NEAR(result.data_ptr<double>()[0], 1.5, 1e-10);
|
||||
ASSERT_NEAR(result.data_ptr<double>()[1], 2.5, 1e-10);
|
||||
ASSERT_NEAR(result.data_ptr<double>()[2], 3.5, 1e-10);
|
||||
}
|
||||
|
||||
TEST(TestAbsolute, AbsoluteMatchesAbs) {
|
||||
// Test that absolute() returns same result as abs()
|
||||
at::Tensor tensor = at::tensor({-3.0f, 2.0f, -1.0f, 0.0f});
|
||||
at::Tensor result_absolute = tensor.absolute();
|
||||
at::Tensor result_abs = tensor.abs();
|
||||
|
||||
ASSERT_EQ(result_absolute.numel(), result_abs.numel());
|
||||
for (int i = 0; i < result_absolute.numel(); ++i) {
|
||||
ASSERT_NEAR(result_absolute.data_ptr<float>()[i],
|
||||
result_abs.data_ptr<float>()[i],
|
||||
1e-6f);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ======================== any tests ========================
|
||||
|
||||
TEST(TensorAnyTest, AnyNoDim) {
|
||||
at::Tensor t = at::zeros({3, 3}, at::kFloat);
|
||||
t.data_ptr<float>()[4] = 1.0f; // Set one element to non-zero
|
||||
|
||||
at::Tensor result = t.any();
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_TRUE(result.item<bool>());
|
||||
}
|
||||
|
||||
TEST(TensorAnyTest, AnyNoDimAllZero) {
|
||||
at::Tensor t = at::zeros({3, 3}, at::kFloat);
|
||||
|
||||
at::Tensor result = t.any();
|
||||
ASSERT_EQ(result.numel(), 1);
|
||||
ASSERT_FALSE(result.item<bool>());
|
||||
}
|
||||
|
||||
TEST(TensorAnyTest, AnyWithDim) {
|
||||
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
||||
t.data_ptr<float>()[0] = 1.0f; // First row has non-zero
|
||||
|
||||
at::Tensor result = t.any(0);
|
||||
ASSERT_EQ(result.numel(), 3);
|
||||
}
|
||||
|
||||
TEST(TensorAnyTest, AnyWithDimKeepdim) {
|
||||
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
||||
t.data_ptr<float>()[0] = 1.0f;
|
||||
|
||||
at::Tensor result = t.any(0, true);
|
||||
ASSERT_EQ(result.sizes().size(), 2);
|
||||
ASSERT_EQ(result.size(0), 1);
|
||||
ASSERT_EQ(result.size(1), 3);
|
||||
}
|
||||
|
||||
TEST(TensorAnyTest, AnyWithOptionalDim) {
|
||||
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
||||
t.data_ptr<float>()[3] = 1.0f; // Second row has non-zero
|
||||
|
||||
at::Tensor result = t.any(at::OptionalIntArrayRef(1));
|
||||
ASSERT_EQ(result.numel(), 2);
|
||||
}
|
||||
|
||||
TEST(TensorAnyTest, AnyInt) {
|
||||
at::Tensor t = at::zeros({2, 3}, at::kInt);
|
||||
t.data_ptr<int>()[0] = 1;
|
||||
|
||||
at::Tensor result = t.any(1);
|
||||
ASSERT_EQ(result.numel(), 2);
|
||||
}
|
||||
|
||||
TEST(TensorAnyTest, AnyBool) {
|
||||
at::Tensor t = at::zeros({3, 3}, at::kBool);
|
||||
t.data_ptr<bool>()[4] = true;
|
||||
|
||||
at::Tensor result = t.any();
|
||||
ASSERT_TRUE(result.item<bool>());
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/common/macros.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_stride_kernel);
|
||||
|
||||
namespace {
|
||||
|
||||
class TensorAsStridedTest : public ::testing::Test {};
|
||||
|
||||
} // namespace
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedBasic) {
|
||||
// shape {2,3}, stride {3,1}: [[0,1,2],[3,4,5]]
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor result = t.as_strided({2, 3}, {3, 1});
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 5.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedWithOffset) {
|
||||
// offset=2: [[2,3,4],[5,6,7]]
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor result = t.as_strided({2, 3}, {3, 1}, 2);
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[5], 7.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedWithDifferentStrides) {
|
||||
// shape {4,2}, stride {2,1}: [[0,1],[2,3],[4,5],[6,7]]
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor result = t.as_strided({4, 2}, {2, 1});
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({4, 2}));
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[7], 7.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedInplace) {
|
||||
// inplace: shape {12} -> {2,6}
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
float* original_data_ptr = t.data_ptr<float>();
|
||||
|
||||
t.as_strided_({2, 6}, {6, 1});
|
||||
|
||||
ASSERT_EQ(t.sizes(), c10::IntArrayRef({2, 6}));
|
||||
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[11], 11.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedInplaceWithOffset) {
|
||||
// inplace with offset=1: [[1,2,3],[4,5,6]]
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
float* original_data_ptr = t.data_ptr<float>();
|
||||
|
||||
t.as_strided_({2, 3}, {3, 1}, 1);
|
||||
|
||||
ASSERT_EQ(t.sizes(), c10::IntArrayRef({2, 3}));
|
||||
ASSERT_NE(t.data_ptr<float>(), original_data_ptr);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 6.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedInplaceModifiesView) {
|
||||
// Modify view, verify original is affected
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor view = t.as_strided({2, 3}, {3, 1});
|
||||
|
||||
view.data_ptr<float>()[0] = 99.0f;
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[0], 99.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedScatterBasic) {
|
||||
// Scatter 2x3 99s into t: [[99,99,99],[99,99,99],...]]
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor src = at::full({2, 3}, 99.0f, at::kFloat);
|
||||
at::Tensor result = t.as_strided_scatter(src, {2, 3}, {3, 1});
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({12}));
|
||||
float* data = result.data_ptr<float>();
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
ASSERT_FLOAT_EQ(data[i], 99.0f);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedScatterOriginalUnchanged) {
|
||||
// Scatter returns new tensor, original unchanged
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor src = at::full({2, 3}, 99.0f, at::kFloat);
|
||||
at::Tensor result = t.as_strided_scatter(src, {2, 3}, {3, 1});
|
||||
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[0], 0.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedScatterWithOffset) {
|
||||
// Scatter with offset=2: [[88,88],[88,88]]
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
at::Tensor src = at::full({2, 2}, 88.0f, at::kFloat);
|
||||
at::Tensor result = t.as_strided_scatter(src, {2, 2}, {2, 1}, 2);
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({12}));
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[2], 88.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 88.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedTranspose) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
// Transpose: shape {2,3} -> {3,2}, stride {1,2}
|
||||
// [[0,1,2],[3,4,5]] -> [[0,3],[1,4],[2,5]]
|
||||
at::Tensor t = at::arange(6, at::kFloat).view({2, 3});
|
||||
at::Tensor result = t.as_strided({3, 2}, {1, 2});
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({3, 2}));
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 5.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorAsStridedTest, AsStridedContiguous) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
at::Tensor t = at::arange(12, at::kFloat);
|
||||
|
||||
// Contiguous: {2,6}, stride {6,1}
|
||||
at::Tensor contig = t.as_strided({2, 6}, {6, 1});
|
||||
ASSERT_TRUE(contig.is_contiguous());
|
||||
|
||||
// Non-contiguous: {3,2}, stride {1,3}
|
||||
at::Tensor non_contig = t.as_strided({3, 2}, {1, 3});
|
||||
ASSERT_FALSE(non_contig.is_contiguous());
|
||||
}
|
||||
@@ -0,0 +1,394 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include <ATen/ops/detach.h>
|
||||
#include <ATen/ops/reciprocal.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// Test detach member function: tensor.detach()
|
||||
TEST(TestDetach, MemberFunction) {
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Detach creates a new tensor that shares data but has no autograd history
|
||||
at::Tensor detached = tensor.detach();
|
||||
|
||||
ASSERT_EQ(detached.sizes(), tensor.sizes());
|
||||
ASSERT_EQ(detached.numel(), tensor.numel());
|
||||
ASSERT_EQ(detached.dtype(), tensor.dtype());
|
||||
|
||||
// Both tensors should share the same data
|
||||
float* original_ptr = tensor.data_ptr<float>();
|
||||
float* detached_ptr = detached.data_ptr<float>();
|
||||
ASSERT_EQ(original_ptr, detached_ptr);
|
||||
}
|
||||
|
||||
// Test detach free function: at::detach(tensor)
|
||||
TEST(TestDetach, FreeFunction) {
|
||||
at::Tensor tensor = at::ones({3, 4}, at::kFloat);
|
||||
at::Tensor detached = at::detach(tensor);
|
||||
|
||||
ASSERT_EQ(detached.sizes(), tensor.sizes());
|
||||
ASSERT_EQ(detached.numel(), tensor.numel());
|
||||
|
||||
// Verify data is shared
|
||||
float* original_ptr = tensor.data_ptr<float>();
|
||||
float* detached_ptr = detached.data_ptr<float>();
|
||||
ASSERT_EQ(original_ptr, detached_ptr);
|
||||
}
|
||||
|
||||
// Test that both methods produce identical results (shared implementation)
|
||||
TEST(TestDetach, SharedImplementation) {
|
||||
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
||||
|
||||
// Call both detach methods
|
||||
at::Tensor detached_member = tensor.detach();
|
||||
at::Tensor detached_free = at::detach(tensor);
|
||||
|
||||
// Both should have the same properties
|
||||
ASSERT_EQ(detached_member.sizes(), detached_free.sizes());
|
||||
ASSERT_EQ(detached_member.numel(), detached_free.numel());
|
||||
ASSERT_EQ(detached_member.dtype(), detached_free.dtype());
|
||||
|
||||
// All three should share the same data
|
||||
float* original_ptr = tensor.data_ptr<float>();
|
||||
float* member_ptr = detached_member.data_ptr<float>();
|
||||
float* free_ptr = detached_free.data_ptr<float>();
|
||||
|
||||
ASSERT_EQ(original_ptr, member_ptr);
|
||||
ASSERT_EQ(original_ptr, free_ptr);
|
||||
}
|
||||
|
||||
// Test detach_ in-place member function: tensor.detach_()
|
||||
TEST(TestDetach, InplaceMemberFunction) {
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
void* original_ptr = tensor.data_ptr();
|
||||
|
||||
// detach_() modifies the tensor in-place
|
||||
at::Tensor& result = tensor.detach_();
|
||||
|
||||
// Should return reference to the same tensor
|
||||
ASSERT_EQ(&result, &tensor);
|
||||
ASSERT_EQ(result.data_ptr(), original_ptr);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
||||
}
|
||||
|
||||
// Test detach preserves data values
|
||||
TEST(TestDetach, PreservesData) {
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
float* data = tensor.data_ptr<float>();
|
||||
data[0] = 1.0f;
|
||||
data[1] = 2.0f;
|
||||
data[2] = 3.0f;
|
||||
data[3] = 4.0f;
|
||||
data[4] = 5.0f;
|
||||
data[5] = 6.0f;
|
||||
|
||||
at::Tensor detached = tensor.detach();
|
||||
|
||||
// Verify data is preserved
|
||||
float* detached_data = detached.data_ptr<float>();
|
||||
ASSERT_EQ(detached_data[0], 1.0f);
|
||||
ASSERT_EQ(detached_data[1], 2.0f);
|
||||
ASSERT_EQ(detached_data[2], 3.0f);
|
||||
ASSERT_EQ(detached_data[3], 4.0f);
|
||||
ASSERT_EQ(detached_data[4], 5.0f);
|
||||
ASSERT_EQ(detached_data[5], 6.0f);
|
||||
}
|
||||
|
||||
// Test detach with different dtypes
|
||||
TEST(TestDetach, DifferentDtypes) {
|
||||
// Float32
|
||||
at::Tensor float_tensor = at::ones({2, 3}, at::kFloat);
|
||||
at::Tensor float_detached = float_tensor.detach();
|
||||
ASSERT_EQ(float_detached.dtype(), at::kFloat);
|
||||
ASSERT_EQ(float_detached.sizes(), float_tensor.sizes());
|
||||
|
||||
// Float64
|
||||
at::Tensor double_tensor = at::ones({2, 3}, at::kDouble);
|
||||
at::Tensor double_detached = double_tensor.detach();
|
||||
ASSERT_EQ(double_detached.dtype(), at::kDouble);
|
||||
ASSERT_EQ(double_detached.sizes(), double_tensor.sizes());
|
||||
|
||||
// Int32
|
||||
at::Tensor int_tensor = at::ones({2, 3}, at::kInt);
|
||||
at::Tensor int_detached = int_tensor.detach();
|
||||
ASSERT_EQ(int_detached.dtype(), at::kInt);
|
||||
ASSERT_EQ(int_detached.sizes(), int_tensor.sizes());
|
||||
|
||||
// Int64
|
||||
at::Tensor long_tensor = at::ones({2, 3}, at::kLong);
|
||||
at::Tensor long_detached = long_tensor.detach();
|
||||
ASSERT_EQ(long_detached.dtype(), at::kLong);
|
||||
ASSERT_EQ(long_detached.sizes(), long_tensor.sizes());
|
||||
}
|
||||
|
||||
// Test detach with various shapes
|
||||
TEST(TestDetach, VariousShapes) {
|
||||
// 1D tensor
|
||||
at::Tensor tensor_1d = at::ones({10}, at::kFloat);
|
||||
at::Tensor detached_1d = tensor_1d.detach();
|
||||
ASSERT_EQ(detached_1d.sizes(), c10::IntArrayRef({10}));
|
||||
|
||||
// 2D tensor
|
||||
at::Tensor tensor_2d = at::ones({3, 4}, at::kFloat);
|
||||
at::Tensor detached_2d = tensor_2d.detach();
|
||||
ASSERT_EQ(detached_2d.sizes(), c10::IntArrayRef({3, 4}));
|
||||
|
||||
// 3D tensor
|
||||
at::Tensor tensor_3d = at::ones({2, 3, 4}, at::kFloat);
|
||||
at::Tensor detached_3d = tensor_3d.detach();
|
||||
ASSERT_EQ(detached_3d.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
|
||||
// 4D tensor
|
||||
at::Tensor tensor_4d = at::ones({2, 3, 4, 5}, at::kFloat);
|
||||
at::Tensor detached_4d = tensor_4d.detach();
|
||||
ASSERT_EQ(detached_4d.sizes(), c10::IntArrayRef({2, 3, 4, 5}));
|
||||
}
|
||||
|
||||
// Test modifications affect both tensors (shared data)
|
||||
TEST(TestDetach, SharedDataModification) {
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
at::Tensor detached = tensor.detach();
|
||||
|
||||
// Modify original tensor
|
||||
float* tensor_data = tensor.data_ptr<float>();
|
||||
tensor_data[0] = 99.0f;
|
||||
|
||||
// Check that detached tensor sees the change
|
||||
float* detached_data = detached.data_ptr<float>();
|
||||
ASSERT_EQ(detached_data[0], 99.0f);
|
||||
|
||||
// Modify detached tensor
|
||||
detached_data[1] = 88.0f;
|
||||
|
||||
// Check that original tensor sees the change
|
||||
ASSERT_EQ(tensor_data[1], 88.0f);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Reciprocal Tests
|
||||
// ============================================================================
|
||||
|
||||
// Test reciprocal member function: tensor.reciprocal()
|
||||
TEST(TestReciprocal, MemberFunction) {
|
||||
at::Tensor tensor = at::full({2, 3}, 2.0f, at::kFloat);
|
||||
at::Tensor result = tensor.reciprocal();
|
||||
|
||||
ASSERT_EQ(result.sizes(), tensor.sizes());
|
||||
ASSERT_EQ(result.numel(), tensor.numel());
|
||||
|
||||
// Verify reciprocal calculation: 1/2 = 0.5
|
||||
float* result_data = result.data_ptr<float>();
|
||||
for (int i = 0; i < result.numel(); i++) {
|
||||
ASSERT_NEAR(result_data[i], 0.5f, 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
// Test reciprocal free function: at::reciprocal(tensor)
|
||||
TEST(TestReciprocal, FreeFunction) {
|
||||
at::Tensor tensor = at::full({3, 4}, 4.0f, at::kFloat);
|
||||
at::Tensor result = at::reciprocal(tensor);
|
||||
|
||||
ASSERT_EQ(result.sizes(), tensor.sizes());
|
||||
ASSERT_EQ(result.numel(), tensor.numel());
|
||||
|
||||
// Verify reciprocal calculation: 1/4 = 0.25
|
||||
float* result_data = result.data_ptr<float>();
|
||||
for (int i = 0; i < result.numel(); i++) {
|
||||
ASSERT_NEAR(result_data[i], 0.25f, 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
// Test that both methods produce identical results (shared implementation)
|
||||
TEST(TestReciprocal, SharedImplementation) {
|
||||
at::Tensor tensor = at::full({2, 3, 4}, 5.0f, at::kFloat);
|
||||
|
||||
// Call both reciprocal methods
|
||||
at::Tensor result_member = tensor.reciprocal();
|
||||
at::Tensor result_free = at::reciprocal(tensor);
|
||||
|
||||
// Both should have the same shape and values
|
||||
ASSERT_EQ(result_member.sizes(), result_free.sizes());
|
||||
ASSERT_EQ(result_member.numel(), result_free.numel());
|
||||
|
||||
// Verify both produce same values: 1/5 = 0.2
|
||||
float* member_data = result_member.data_ptr<float>();
|
||||
float* free_data = result_free.data_ptr<float>();
|
||||
for (int i = 0; i < result_member.numel(); i++) {
|
||||
ASSERT_NEAR(member_data[i], 0.2f, 1e-6);
|
||||
ASSERT_NEAR(free_data[i], 0.2f, 1e-6);
|
||||
ASSERT_EQ(member_data[i], free_data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// Test reciprocal_ in-place member function: tensor.reciprocal_()
|
||||
TEST(TestReciprocal, InplaceMemberFunction) {
|
||||
at::Tensor tensor = at::full({2, 3}, 2.0f, at::kFloat);
|
||||
void* original_ptr = tensor.data_ptr();
|
||||
|
||||
// reciprocal_() modifies the tensor in-place
|
||||
at::Tensor& result = tensor.reciprocal_();
|
||||
|
||||
// Should return reference to the same tensor
|
||||
ASSERT_EQ(&result, &tensor);
|
||||
ASSERT_EQ(result.data_ptr(), original_ptr);
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
||||
|
||||
// Verify reciprocal calculation: 1/2 = 0.5
|
||||
float* result_data = result.data_ptr<float>();
|
||||
for (int i = 0; i < result.numel(); i++) {
|
||||
ASSERT_NEAR(result_data[i], 0.5f, 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
// Test reciprocal with various input values
|
||||
TEST(TestReciprocal, VariousValues) {
|
||||
at::Tensor tensor = at::ones({5}, at::kFloat);
|
||||
float* data = tensor.data_ptr<float>();
|
||||
data[0] = 1.0f;
|
||||
data[1] = 2.0f;
|
||||
data[2] = 4.0f;
|
||||
data[3] = 0.5f;
|
||||
data[4] = 10.0f;
|
||||
|
||||
at::Tensor result = tensor.reciprocal();
|
||||
float* result_data = result.data_ptr<float>();
|
||||
|
||||
// Verify reciprocals
|
||||
ASSERT_NEAR(result_data[0], 1.0f, 1e-6); // 1/1 = 1
|
||||
ASSERT_NEAR(result_data[1], 0.5f, 1e-6); // 1/2 = 0.5
|
||||
ASSERT_NEAR(result_data[2], 0.25f, 1e-6); // 1/4 = 0.25
|
||||
ASSERT_NEAR(result_data[3], 2.0f, 1e-6); // 1/0.5 = 2
|
||||
ASSERT_NEAR(result_data[4], 0.1f, 1e-6); // 1/10 = 0.1
|
||||
}
|
||||
|
||||
// Test reciprocal with different dtypes
|
||||
TEST(TestReciprocal, DifferentDtypes) {
|
||||
// Float32
|
||||
at::Tensor float_tensor = at::full({2, 3}, 2.0f, at::kFloat);
|
||||
at::Tensor float_result = float_tensor.reciprocal();
|
||||
ASSERT_EQ(float_result.dtype(), at::kFloat);
|
||||
float* float_data = float_result.data_ptr<float>();
|
||||
ASSERT_NEAR(float_data[0], 0.5f, 1e-6);
|
||||
|
||||
// Float64
|
||||
at::Tensor double_tensor = at::full({2, 3}, 2.0, at::kDouble);
|
||||
at::Tensor double_result = double_tensor.reciprocal();
|
||||
ASSERT_EQ(double_result.dtype(), at::kDouble);
|
||||
double* double_data = double_result.data_ptr<double>();
|
||||
ASSERT_NEAR(double_data[0], 0.5, 1e-10);
|
||||
}
|
||||
|
||||
// Test reciprocal with various shapes
|
||||
TEST(TestReciprocal, VariousShapes) {
|
||||
// 1D tensor
|
||||
at::Tensor tensor_1d = at::full({10}, 2.0f, at::kFloat);
|
||||
at::Tensor result_1d = tensor_1d.reciprocal();
|
||||
ASSERT_EQ(result_1d.sizes(), c10::IntArrayRef({10}));
|
||||
ASSERT_NEAR(result_1d.data_ptr<float>()[0], 0.5f, 1e-6);
|
||||
|
||||
// 2D tensor
|
||||
at::Tensor tensor_2d = at::full({3, 4}, 2.0f, at::kFloat);
|
||||
at::Tensor result_2d = tensor_2d.reciprocal();
|
||||
ASSERT_EQ(result_2d.sizes(), c10::IntArrayRef({3, 4}));
|
||||
ASSERT_NEAR(result_2d.data_ptr<float>()[0], 0.5f, 1e-6);
|
||||
|
||||
// 3D tensor
|
||||
at::Tensor tensor_3d = at::full({2, 3, 4}, 2.0f, at::kFloat);
|
||||
at::Tensor result_3d = tensor_3d.reciprocal();
|
||||
ASSERT_EQ(result_3d.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
ASSERT_NEAR(result_3d.data_ptr<float>()[0], 0.5f, 1e-6);
|
||||
|
||||
// 4D tensor
|
||||
at::Tensor tensor_4d = at::full({2, 3, 4, 5}, 2.0f, at::kFloat);
|
||||
at::Tensor result_4d = tensor_4d.reciprocal();
|
||||
ASSERT_EQ(result_4d.sizes(), c10::IntArrayRef({2, 3, 4, 5}));
|
||||
ASSERT_NEAR(result_4d.data_ptr<float>()[0], 0.5f, 1e-6);
|
||||
}
|
||||
|
||||
// Test reciprocal_ modifies original tensor
|
||||
TEST(TestReciprocal, InplaceModifiesOriginal) {
|
||||
at::Tensor tensor = at::full({3, 3}, 4.0f, at::kFloat);
|
||||
|
||||
// Store original data pointer
|
||||
void* original_ptr = tensor.data_ptr();
|
||||
|
||||
// Call in-place reciprocal
|
||||
tensor.reciprocal_();
|
||||
|
||||
// Same memory location
|
||||
ASSERT_EQ(tensor.data_ptr(), original_ptr);
|
||||
|
||||
// Values should be modified: 1/4 = 0.25
|
||||
float* data = tensor.data_ptr<float>();
|
||||
for (int i = 0; i < tensor.numel(); i++) {
|
||||
ASSERT_NEAR(data[i], 0.25f, 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
// Test reciprocal creates new tensor (non-inplace)
|
||||
TEST(TestReciprocal, CreatesNewTensor) {
|
||||
at::Tensor tensor = at::full({2, 3}, 2.0f, at::kFloat);
|
||||
void* original_ptr = tensor.data_ptr();
|
||||
|
||||
// Non-inplace reciprocal should create new tensor
|
||||
at::Tensor result = tensor.reciprocal();
|
||||
|
||||
// Different memory location
|
||||
ASSERT_NE(result.data_ptr(), original_ptr);
|
||||
|
||||
// Original tensor unchanged
|
||||
float* original_data = tensor.data_ptr<float>();
|
||||
ASSERT_NEAR(original_data[0], 2.0f, 1e-6);
|
||||
|
||||
// Result has reciprocal values
|
||||
float* result_data = result.data_ptr<float>();
|
||||
ASSERT_NEAR(result_data[0], 0.5f, 1e-6);
|
||||
}
|
||||
|
||||
// Test reciprocal with negative values
|
||||
TEST(TestReciprocal, NegativeValues) {
|
||||
at::Tensor tensor = at::ones({4}, at::kFloat);
|
||||
float* data = tensor.data_ptr<float>();
|
||||
data[0] = -1.0f;
|
||||
data[1] = -2.0f;
|
||||
data[2] = -0.5f;
|
||||
data[3] = -4.0f;
|
||||
|
||||
at::Tensor result = tensor.reciprocal();
|
||||
float* result_data = result.data_ptr<float>();
|
||||
|
||||
// Verify reciprocals of negative numbers
|
||||
ASSERT_NEAR(result_data[0], -1.0f, 1e-6); // 1/(-1) = -1
|
||||
ASSERT_NEAR(result_data[1], -0.5f, 1e-6); // 1/(-2) = -0.5
|
||||
ASSERT_NEAR(result_data[2], -2.0f, 1e-6); // 1/(-0.5) = -2
|
||||
ASSERT_NEAR(result_data[3], -0.25f, 1e-6); // 1/(-4) = -0.25
|
||||
}
|
||||
@@ -0,0 +1,518 @@
|
||||
// Copyright (c) 2025 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#ifdef PADDLE_WITH_XPU
|
||||
#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/common/macros.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_stride_kernel);
|
||||
|
||||
TEST(TensorBaseTest, DataPtrAPIs) {
|
||||
// Test data_ptr() and const_data_ptr() APIs
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Test void* data_ptr()
|
||||
void* void_ptr = tensor.data_ptr();
|
||||
ASSERT_NE(void_ptr, nullptr);
|
||||
|
||||
// Test typed data_ptr<T>()
|
||||
float* float_ptr = tensor.data_ptr<float>();
|
||||
ASSERT_NE(float_ptr, nullptr);
|
||||
ASSERT_EQ(float_ptr, void_ptr);
|
||||
|
||||
// Test const_data_ptr()
|
||||
const float* const_float_ptr = tensor.const_data_ptr<float>();
|
||||
ASSERT_NE(const_float_ptr, nullptr);
|
||||
ASSERT_EQ(const_float_ptr, float_ptr);
|
||||
|
||||
// Test mutable_data_ptr()
|
||||
void* mutable_ptr = tensor.mutable_data_ptr();
|
||||
ASSERT_NE(mutable_ptr, nullptr);
|
||||
ASSERT_EQ(mutable_ptr, void_ptr);
|
||||
}
|
||||
TEST(TensorBaseTest, TypeDeviceAPIs) {
|
||||
// Test type and device related APIs
|
||||
at::TensorBase cpu_tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Test dtype()/scalar_type()
|
||||
ASSERT_EQ(cpu_tensor.dtype(), at::kFloat);
|
||||
ASSERT_EQ(cpu_tensor.scalar_type(), at::kFloat);
|
||||
|
||||
// Test device()
|
||||
ASSERT_EQ(cpu_tensor.device().type(), at::DeviceType::CPU);
|
||||
|
||||
// Test get_device()
|
||||
ASSERT_EQ(cpu_tensor.get_device(), -1); // CPU device index is -1
|
||||
|
||||
// Test is_cpu()/is_cuda()
|
||||
ASSERT_TRUE(cpu_tensor.is_cpu());
|
||||
ASSERT_FALSE(cpu_tensor.is_cuda());
|
||||
|
||||
// Test options()
|
||||
auto options = cpu_tensor.options();
|
||||
ASSERT_EQ(options.device().type(), at::DeviceType::CPU);
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, ModifyOperationAPIs) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
// Test modify operation related APIs
|
||||
at::Tensor tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Test is_contiguous()
|
||||
ASSERT_TRUE(tensor.is_contiguous());
|
||||
|
||||
// Test is_contiguous_or_false()
|
||||
ASSERT_TRUE(tensor.is_contiguous_or_false());
|
||||
|
||||
// Test fill_()
|
||||
tensor.fill_(2.0);
|
||||
float* data = tensor.data_ptr<float>();
|
||||
for (int i = 0; i < tensor.numel(); i++) {
|
||||
ASSERT_EQ(data[i], 2.0f);
|
||||
}
|
||||
|
||||
// Test zero_()
|
||||
tensor.zero_();
|
||||
for (int i = 0; i < tensor.numel(); i++) {
|
||||
ASSERT_EQ(data[i], 0.0f);
|
||||
}
|
||||
|
||||
// Test copy_()
|
||||
at::Tensor src = at::ones({2, 3}, at::kFloat);
|
||||
tensor.copy_(src);
|
||||
for (int i = 0; i < tensor.numel(); i++) {
|
||||
ASSERT_EQ(data[i], 1.0f);
|
||||
}
|
||||
|
||||
// Test view()
|
||||
at::TensorBase viewed = tensor.view({6});
|
||||
ASSERT_EQ(viewed.sizes(), std::vector<int64_t>{6});
|
||||
ASSERT_EQ(viewed.strides(), std::vector<int64_t>{1});
|
||||
}
|
||||
|
||||
TEST(tensor_clone_test, BasicClone) {
|
||||
at::Tensor a = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
at::Tensor b = a.clone();
|
||||
|
||||
ASSERT_EQ(a.sizes(), b.sizes());
|
||||
ASSERT_EQ(a.dtype(), b.dtype());
|
||||
ASSERT_EQ(a.device().type(), b.device().type());
|
||||
}
|
||||
|
||||
TEST(compat_basic_test, BasicCase) {
|
||||
at::Tensor a =
|
||||
at::ones({2, 3}, at::TensorOptions().dtype(at::kFloat).device(at::kCPU));
|
||||
at::Tensor b = at::full({2, 3}, 2, at::kFloat);
|
||||
double c = 10;
|
||||
|
||||
TORCH_CHECK(a.sizes() == b.sizes());
|
||||
TORCH_CHECK(a.dtype() == at::kFloat);
|
||||
TORCH_CHECK(b.dtype() == at::kFloat);
|
||||
TORCH_INTERNAL_ASSERT(a.device().type() == at::DeviceType::CPU);
|
||||
TORCH_INTERNAL_ASSERT(b.device().type() == at::DeviceType::CPU);
|
||||
at::Tensor a_contig = a.contiguous();
|
||||
at::Tensor b_contig = b.contiguous();
|
||||
at::Tensor result = at::empty(a_contig.sizes(), a_contig.options());
|
||||
const float* a_ptr = a_contig.data_ptr<float>();
|
||||
const float* b_ptr = b_contig.data_ptr<float>();
|
||||
float* result_ptr = result.data_ptr<float>();
|
||||
for (int64_t i = 0; i < a_contig.numel(); i++) {
|
||||
result_ptr[i] = a_ptr[i] * b_ptr[i] + c;
|
||||
}
|
||||
// Show result
|
||||
for (int64_t i = 0; i < a_contig.numel(); i++) {
|
||||
std::cout << "Result[" << i << "] = " << a_ptr[i] * b_ptr[i] + c
|
||||
<< std::endl;
|
||||
ASSERT_EQ(result_ptr[i], 12);
|
||||
}
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
|
||||
{
|
||||
// for test empty_cuda:
|
||||
at::Tensor bb =
|
||||
at::detail::empty_cuda(12, at::kFloat, at::kCUDA, std::nullopt);
|
||||
|
||||
// for test sizoof(at::Half):
|
||||
std::cout << sizeof(at::Half) << std::endl;
|
||||
at::Tensor num_non_exiting_ctas = at::empty(
|
||||
{}, at::TensorOptions().device(a.device()).dtype(at::ScalarType::Int));
|
||||
}
|
||||
{
|
||||
std::vector<int64_t> shape = {2, 3, 4, 5};
|
||||
size_t size_ =
|
||||
c10::elementSize(at::ScalarType::Float) * c10::multiply_integers(shape);
|
||||
std::cout << "multiply_integers out: " << size_ << std::endl;
|
||||
}
|
||||
{
|
||||
std::vector<int> shape = {2, 3, 4, 5};
|
||||
size_t size_ =
|
||||
c10::elementSize(at::ScalarType::Float) * c10::sum_integers(shape);
|
||||
std::cout << "sum_integers out: " << size_ << std::endl;
|
||||
}
|
||||
{
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
std::cout << "stream num: " << stream.stream() << std::endl;
|
||||
at::cuda::stream_synchronize(stream);
|
||||
at::Tensor bb =
|
||||
at::detail::empty_cuda(12, at::kFloat, at::kCUDA, std::nullopt);
|
||||
}
|
||||
{
|
||||
at::Tensor a = at::ones(
|
||||
{2, 3}, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
std::cout << "a.device() is at::kCUDA: " << (a.device().type() == at::kCUDA)
|
||||
<< std::endl;
|
||||
const c10::cuda::CUDAGuard device_guard(a.device());
|
||||
std::cout << "device_guard is at::kCUDA: "
|
||||
<< (device_guard.current_device().type() == at::kCUDA)
|
||||
<< std::endl;
|
||||
const c10::cuda::OptionalCUDAGuard device_guard_opt(a.device());
|
||||
std::cout << "device_guard is at::kCUDA: "
|
||||
<< (device_guard_opt.current_device().value().type() == at::kCUDA)
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
{
|
||||
std::cout << "num_tokens_per_rank.device() is at::kCUDA: " << std::endl;
|
||||
// for test empty:
|
||||
auto num_tokens_per_rank =
|
||||
torch::empty({3},
|
||||
dtype(torch::kInt32).device(torch::kCUDA),
|
||||
c10::MemoryFormat::Contiguous);
|
||||
std::cout << "num_tokens_per_rank.device() is at::kCUDA: "
|
||||
<< (num_tokens_per_rank.device().type() == at::kCUDA)
|
||||
<< std::endl;
|
||||
}
|
||||
{
|
||||
auto num_tokens_per_rank = torch::empty(
|
||||
{3}, dtype(torch::kInt32).device(torch::kCUDA), std::nullopt);
|
||||
std::cout << "num_tokens_per_rank.device() is at::kCUDA: "
|
||||
<< (num_tokens_per_rank.device().type() == at::kCUDA)
|
||||
<< std::endl;
|
||||
}
|
||||
#endif
|
||||
{
|
||||
int a = 10, b = 20, c = 30;
|
||||
int* p[] = {&a, &b, &c}; // int* array[3]
|
||||
int** pp = p;
|
||||
|
||||
torch::Tensor t =
|
||||
torch::from_blob(pp, {3}, torch::TensorOptions().dtype(torch::kInt64));
|
||||
|
||||
// Get original int**
|
||||
int** restored = reinterpret_cast<int**>(t.data_ptr<int64_t>());
|
||||
std::cout << *restored[0] << ", " << *restored[1] << ", " << *restored[2]
|
||||
<< std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TestDevice, DeviceAPIsOnCUDA) {
|
||||
// Test device related APIs on CUDA if available
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
if (at::cuda::is_available()) {
|
||||
at::TensorBase cuda_tensor = at::ones(
|
||||
{2, 3}, c10::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
|
||||
// Test device()
|
||||
ASSERT_EQ(cuda_tensor.device().type(), at::DeviceType::CUDA);
|
||||
|
||||
// Test get_device()
|
||||
ASSERT_EQ(cuda_tensor.get_device(), 0); // Assuming single GPU with index 0
|
||||
|
||||
// Test is_cpu()/is_cuda()
|
||||
ASSERT_FALSE(cuda_tensor.is_cpu());
|
||||
ASSERT_TRUE(cuda_tensor.is_cuda());
|
||||
|
||||
// Test options()
|
||||
auto options = cuda_tensor.options();
|
||||
ASSERT_EQ(options.device().type(), at::DeviceType::CUDA);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
TEST(TestDevice, DeviceAPIsOnCPU) {
|
||||
// Test device related APIs on CPU
|
||||
at::TensorBase cpu_tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Test device()
|
||||
ASSERT_EQ(cpu_tensor.device().type(), at::DeviceType::CPU);
|
||||
|
||||
// Test is_cpu()/is_cuda()
|
||||
ASSERT_TRUE(cpu_tensor.is_cpu());
|
||||
ASSERT_FALSE(cpu_tensor.is_cuda());
|
||||
|
||||
// Test options()
|
||||
auto options = cpu_tensor.options();
|
||||
ASSERT_EQ(options.device().type(), at::DeviceType::CPU);
|
||||
}
|
||||
|
||||
TEST(TestTranspose, TransposeAPI) {
|
||||
at::Tensor a = at::ones({4, 5, 6, 7, 8}, at::kFloat);
|
||||
at::Tensor b = a.transpose(2, 3);
|
||||
ASSERT_EQ(b.sizes(), c10::IntArrayRef({4, 5, 7, 6, 8}));
|
||||
}
|
||||
|
||||
TEST(TestSize, SizeNegativeIndex) {
|
||||
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
||||
ASSERT_EQ(tensor.size(-1), 5);
|
||||
ASSERT_EQ(tensor.size(-2), 4);
|
||||
ASSERT_EQ(tensor.size(-3), 3);
|
||||
ASSERT_EQ(tensor.size(-4), 2);
|
||||
}
|
||||
|
||||
TEST(TestTensorOperators, SubScriptOperator) {
|
||||
const int M = 3;
|
||||
const int N = 4;
|
||||
const int K = 5;
|
||||
|
||||
at::Tensor tensor = at::arange(M * N * K, at::kFloat).reshape({M, N, K});
|
||||
|
||||
// Check tensor[0]
|
||||
at::Tensor tensor_0 = tensor[0];
|
||||
for (int i = 0; i < N * K; ++i) {
|
||||
ASSERT_EQ(tensor_0.data_ptr<float>()[i], static_cast<float>(i));
|
||||
}
|
||||
|
||||
// Check tensor[1]
|
||||
at::Tensor tensor_1 = tensor[1];
|
||||
int offset = N * K;
|
||||
for (int i = 0; i < N * K; ++i) {
|
||||
ASSERT_EQ(tensor_1.data_ptr<float>()[i], static_cast<float>(i + offset));
|
||||
}
|
||||
|
||||
// Check tensor[2]
|
||||
at::Tensor tensor_2 = tensor[2];
|
||||
offset = 2 * N * K;
|
||||
for (int i = 0; i < N * K; ++i) {
|
||||
ASSERT_EQ(tensor_2.data_ptr<float>()[i], static_cast<float>(i + offset));
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, LayoutAPI) {
|
||||
// Test layout() API for strided tensors
|
||||
at::TensorBase tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Default tensor should have Strided layout
|
||||
ASSERT_EQ(tensor.layout(), c10::kStrided);
|
||||
|
||||
// Test layout output stream operator
|
||||
std::ostringstream oss;
|
||||
oss << tensor.layout();
|
||||
ASSERT_EQ(oss.str(), "Strided");
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, ResetAPI) {
|
||||
// Test reset() API
|
||||
at::TensorBase tensor = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
// Verify tensor is defined before reset
|
||||
ASSERT_TRUE(tensor.defined());
|
||||
ASSERT_NE(tensor.data_ptr(), nullptr);
|
||||
ASSERT_EQ(tensor.numel(), 6);
|
||||
|
||||
// Call reset()
|
||||
tensor.reset();
|
||||
|
||||
// Verify tensor is no longer defined after reset
|
||||
ASSERT_FALSE(tensor.defined());
|
||||
|
||||
// Test reset on already undefined tensor (should not crash)
|
||||
at::TensorBase empty_tensor;
|
||||
ASSERT_FALSE(empty_tensor.defined());
|
||||
empty_tensor.reset();
|
||||
ASSERT_FALSE(empty_tensor.defined());
|
||||
|
||||
// Test reset on tensor after assignment
|
||||
at::TensorBase tensor2 = at::ones({3, 4}, at::kDouble);
|
||||
at::TensorBase tensor3 = tensor2;
|
||||
ASSERT_TRUE(tensor2.defined());
|
||||
ASSERT_TRUE(tensor3.defined());
|
||||
|
||||
tensor2.reset();
|
||||
ASSERT_FALSE(tensor2.defined());
|
||||
ASSERT_TRUE(tensor3.defined()); // tensor3 should still be valid
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, IsNonOverlappingAndDenseAPI) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
// Test is_non_overlapping_and_dense() API
|
||||
|
||||
// Case 1: Contiguous tensor - should be non-overlapping and dense
|
||||
at::TensorBase contiguous_tensor = at::ones({2, 3, 4}, at::kFloat);
|
||||
ASSERT_TRUE(contiguous_tensor.is_contiguous());
|
||||
ASSERT_TRUE(contiguous_tensor.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 2: Scalar tensor (numel == 1) - should be non-overlapping and dense
|
||||
at::TensorBase scalar_tensor = at::ones({}, at::kFloat);
|
||||
ASSERT_EQ(scalar_tensor.numel(), 1);
|
||||
ASSERT_TRUE(scalar_tensor.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 3: Single element tensor - should be non-overlapping and dense
|
||||
at::TensorBase single_element = at::ones({1, 1, 1}, at::kFloat);
|
||||
ASSERT_EQ(single_element.numel(), 1);
|
||||
ASSERT_TRUE(single_element.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 4: Transposed tensor - non-contiguous but still non-overlapping and
|
||||
// dense
|
||||
at::Tensor original = at::ones({3, 4}, at::kFloat);
|
||||
at::Tensor transposed = original.transpose(0, 1);
|
||||
ASSERT_FALSE(transposed.is_contiguous());
|
||||
ASSERT_TRUE(transposed.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 5: Multi-dimensional transpose - still non-overlapping and dense
|
||||
at::Tensor tensor_3d = at::ones({2, 3, 4}, at::kFloat);
|
||||
at::Tensor transposed_3d = tensor_3d.transpose(0, 2);
|
||||
ASSERT_FALSE(transposed_3d.is_contiguous());
|
||||
ASSERT_TRUE(transposed_3d.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 6: Tensor with size-1 dimensions
|
||||
at::TensorBase size_one_dims = at::ones({1, 3, 1, 4}, at::kFloat);
|
||||
ASSERT_TRUE(size_one_dims.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 7: Empty tensor (numel == 0) - should be non-overlapping and dense
|
||||
at::TensorBase empty_tensor = at::ones({0, 3, 4}, at::kFloat);
|
||||
ASSERT_EQ(empty_tensor.numel(), 0);
|
||||
ASSERT_TRUE(empty_tensor.is_non_overlapping_and_dense());
|
||||
|
||||
// Case 8: Permuted tensor - still non-overlapping and dense
|
||||
at::Tensor tensor_4d = at::ones({2, 3, 4, 5}, at::kFloat);
|
||||
at::Tensor permuted = tensor_4d.permute({3, 1, 2, 0});
|
||||
ASSERT_FALSE(permuted.is_contiguous());
|
||||
ASSERT_TRUE(permuted.is_non_overlapping_and_dense());
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, UndefinedAndNonDenseBranchCoverage) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
at::TensorBase undefined;
|
||||
ASSERT_EQ(undefined.toString(), std::string("UndefinedType"));
|
||||
ASSERT_EQ(undefined.data_ptr(), nullptr);
|
||||
ASSERT_FALSE(undefined.has_names());
|
||||
|
||||
at::Tensor non_dense = at::arange(6, at::TensorOptions().dtype(at::kFloat))
|
||||
.as_strided({2, 2}, {4, 1});
|
||||
ASSERT_FALSE(non_dense.is_non_overlapping_and_dense());
|
||||
}
|
||||
|
||||
TEST(TensorBodyTest, ToBackendUnsupportedBranch) {
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
ASSERT_THROW(t.toBackend(static_cast<c10::Backend>(-1)), ::std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorBodyTest, ToBackendCpuBranchCoverage) {
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
at::Tensor cpu_t = t.toBackend(c10::Backend::CPU);
|
||||
|
||||
ASSERT_EQ(cpu_t.device().type(), c10::DeviceType::CPU);
|
||||
ASSERT_TRUE(cpu_t.equal(t));
|
||||
}
|
||||
|
||||
TEST(TensorBodyTest, ToBackendCudaBranchCoverage) {
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
|
||||
try {
|
||||
at::Tensor cuda_t = t.toBackend(c10::Backend::CUDA);
|
||||
ASSERT_EQ(cuda_t.device().type(), c10::DeviceType::CUDA);
|
||||
} catch (const std::exception&) {
|
||||
SUCCEED();
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorBodyTest, ToBackendXpuBranchCoverage) {
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
|
||||
try {
|
||||
at::Tensor xpu_t = t.toBackend(c10::Backend::XPU);
|
||||
ASSERT_EQ(xpu_t.device().type(), c10::DeviceType::XPU);
|
||||
} catch (const std::exception&) {
|
||||
SUCCEED();
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorBodyTest, ToBackendIpuBranchCoverage) {
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
|
||||
try {
|
||||
at::Tensor ipu_t = t.toBackend(c10::Backend::IPU);
|
||||
ASSERT_EQ(ipu_t.device().type(), c10::DeviceType::IPU);
|
||||
} catch (const std::exception&) {
|
||||
SUCCEED();
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef PADDLE_WITH_XPU
|
||||
TEST(TensorBodyTest, ToBackendXpuUsesCurrentDevice) {
|
||||
if (paddle::platform::GetXPUDeviceCount() < 2) {
|
||||
return;
|
||||
}
|
||||
paddle::platform::XPUDeviceGuard guard(1);
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
at::Tensor xpu_t = t.toBackend(c10::Backend::XPU);
|
||||
|
||||
ASSERT_EQ(xpu_t.device().type(), c10::DeviceType::XPU);
|
||||
ASSERT_EQ(xpu_t.device().index(), 1);
|
||||
}
|
||||
#endif
|
||||
|
||||
TEST(TensorBodyTest, MetaUnsupportedBranch) {
|
||||
at::Tensor t = at::ones({1}, at::kFloat);
|
||||
ASSERT_THROW((void)t.meta(), ::std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, ToDeviceAndMemoryFormatUnsupportedBranches) {
|
||||
at::TensorBase base = at::ones({2, 2}, at::kFloat);
|
||||
|
||||
ASSERT_THROW(
|
||||
(void)base.to(at::TensorOptions().device(c10::Device(c10::kCPU))),
|
||||
::std::exception);
|
||||
|
||||
ASSERT_THROW((void)base.to(at::TensorOptions().dtype(at::kFloat),
|
||||
false,
|
||||
false,
|
||||
at::MemoryFormat::Contiguous),
|
||||
::std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorBaseTest, ToDtypeCastsWhenSupported) {
|
||||
at::TensorBase base = at::ones({2, 2}, at::kFloat);
|
||||
at::TensorBase casted = base.to(at::TensorOptions().dtype(at::kDouble));
|
||||
ASSERT_EQ(casted.scalar_type(), at::kDouble);
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ======================== chunk tests ========================
|
||||
|
||||
TEST(TensorChunkTest, ChunkBasic) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(3, 0);
|
||||
|
||||
ASSERT_EQ(chunks.size(), 3);
|
||||
ASSERT_EQ(chunks[0].size(0), 1);
|
||||
ASSERT_EQ(chunks[1].size(0), 1);
|
||||
ASSERT_EQ(chunks[2].size(0), 1);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkDim1) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(2, 1);
|
||||
|
||||
ASSERT_EQ(chunks.size(), 2);
|
||||
ASSERT_EQ(chunks[0].size(1), 2);
|
||||
ASSERT_EQ(chunks[1].size(1), 2);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkUneven) {
|
||||
at::Tensor t = at::arange(10, at::kFloat).reshape({2, 5});
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(3, 1);
|
||||
|
||||
ASSERT_EQ(chunks.size(), 3);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkMoreChunksThanSize) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(5, 0);
|
||||
|
||||
// PyTorch returns at most dim_size non-empty chunks when chunks > dim_size
|
||||
ASSERT_EQ(chunks.size(), 2);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkDefaultDim) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(3);
|
||||
|
||||
ASSERT_EQ(chunks.size(), 3);
|
||||
ASSERT_EQ(chunks[0].size(0), 1);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkIntType) {
|
||||
at::Tensor t = at::arange(12, at::kInt).reshape({3, 4});
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(3, 0);
|
||||
|
||||
ASSERT_EQ(chunks.size(), 3);
|
||||
ASSERT_EQ(chunks[0].dtype(), at::kInt);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkZeroDim) {
|
||||
at::Tensor t = at::zeros({0, 4}, at::kFloat);
|
||||
|
||||
std::vector<at::Tensor> chunks = t.chunk(2, 0);
|
||||
|
||||
// PyTorch returns 'chunks' number of empty tensors when dim_size == 0
|
||||
ASSERT_EQ(chunks.size(), 2);
|
||||
ASSERT_EQ(chunks[0].size(0), 0);
|
||||
ASSERT_EQ(chunks[1].size(0), 0);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkNegativeDim) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
// chunk(-1) should be equivalent to chunk(rank - 1) = chunk(1)
|
||||
std::vector<at::Tensor> chunks_neg = t.chunk(2, -1);
|
||||
std::vector<at::Tensor> chunks_pos = t.chunk(2, 1);
|
||||
|
||||
ASSERT_EQ(chunks_neg.size(), chunks_pos.size());
|
||||
for (size_t i = 0; i < chunks_neg.size(); ++i) {
|
||||
ASSERT_EQ(chunks_neg[i].sizes(), chunks_pos[i].sizes());
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkOutOfRangeDim) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
ASSERT_THROW(t.chunk(2, 2), std::exception); // dim >= rank
|
||||
ASSERT_THROW(t.chunk(2, -3), std::exception); // dim < -rank
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkZeroRankTensor) {
|
||||
at::Tensor t = at::empty({}, at::kFloat); // 0-dim scalar tensor
|
||||
|
||||
ASSERT_THROW(t.chunk(2, 0), std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorChunkTest, ChunkZeroChunks) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
ASSERT_THROW(t.chunk(0, 0), std::exception);
|
||||
ASSERT_THROW(t.chunk(-1, 0), std::exception);
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "test/cpp/prim/init_env_utils.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
namespace {
|
||||
|
||||
class TensorClampTest : public ::testing::Test {
|
||||
protected:
|
||||
static void SetUpTestSuite() { paddle::prim::InitTensorOperants(); }
|
||||
};
|
||||
|
||||
class TensorOperatorIndexTest : public ::testing::Test {
|
||||
protected:
|
||||
static void SetUpTestSuite() { paddle::prim::InitTensorOperants(); }
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
TEST_F(TensorClampTest, ClampWithScalar) {
|
||||
// Create tensor with values [0, 1, 2, 3, 4, 5]
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor result = t.clamp(at::Scalar(1.0), at::Scalar(4.0));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
// Expected: [1, 1, 2, 3, 4, 4]
|
||||
ASSERT_FLOAT_EQ(data[0], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[2], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[3], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[4], 4.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 4.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampWithTensor) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor min_t = at::full({2, 3}, 1.0f, at::kFloat);
|
||||
at::Tensor max_t = at::full({2, 3}, 4.0f, at::kFloat);
|
||||
|
||||
at::Tensor result = t.clamp(::std::optional<at::Tensor>(min_t),
|
||||
::std::optional<at::Tensor>(max_t));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 4.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampInplaceScalar) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
t.clamp_(at::Scalar(2.0), at::Scalar(3.0));
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
// Expected: [2, 2, 2, 3, 3, 3]
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[2], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[3], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[4], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 3.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampInplaceTensor) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor min_t = at::full({2, 3}, 1.0f, at::kFloat);
|
||||
at::Tensor max_t = at::full({2, 3}, 4.0f, at::kFloat);
|
||||
|
||||
t.clamp_(::std::optional<at::Tensor>(min_t),
|
||||
::std::optional<at::Tensor>(max_t));
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 4.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMaxScalar) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor result = t.clamp_max(at::Scalar(3.0));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
// Expected: [0, 1, 2, 3, 3, 3]
|
||||
ASSERT_FLOAT_EQ(data[4], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 3.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMaxTensor) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor max_t = at::full({6}, 3.0f, at::kFloat);
|
||||
at::Tensor result = t.clamp_max(max_t);
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[4], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 3.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMaxInplaceScalar) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
t.clamp_max_(at::Scalar(3.0));
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[4], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 3.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMaxInplaceTensor) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor max_t = at::full({6}, 3.0f, at::kFloat);
|
||||
t.clamp_max_(max_t);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[4], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[5], 3.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinScalar) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor result = t.clamp_min(at::Scalar(2.0));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
// Expected: [2, 2, 2, 3, 4, 5]
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinTensor) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor min_t = at::full({6}, 2.0f, at::kFloat);
|
||||
at::Tensor result = t.clamp_min(min_t);
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinInplaceScalar) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
t.clamp_min_(at::Scalar(2.0));
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinInplaceTensor) {
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor min_t = at::full({6}, 2.0f, at::kFloat);
|
||||
t.clamp_min_(min_t);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 2.0f);
|
||||
}
|
||||
|
||||
// ======================== operator[] tests ========================
|
||||
|
||||
TEST_F(TensorOperatorIndexTest, OperatorIndexBasic) {
|
||||
// Create tensor [[0,1,2],[3,4,5]]
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
// Test operator[](int64_t index) - returns first row
|
||||
at::Tensor result0 = t[0];
|
||||
ASSERT_EQ(result0.numel(), 3); // First row has 3 elements [0,1,2]
|
||||
ASSERT_FLOAT_EQ(result0.data_ptr<float>()[0],
|
||||
0.0f); // First element of the row
|
||||
|
||||
at::Tensor result1 = t[1];
|
||||
ASSERT_EQ(result1.numel(), 3); // Second row has 3 elements [3,4,5]
|
||||
ASSERT_FLOAT_EQ(result1.data_ptr<float>()[0],
|
||||
3.0f); // First element of the row
|
||||
}
|
||||
|
||||
TEST_F(TensorOperatorIndexTest, OperatorIndexOutOfBounds) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
// Test out of bounds index - should throw an exception
|
||||
// The test expects the code to handle this gracefully
|
||||
bool threw_exception = false;
|
||||
try {
|
||||
at::Tensor result = t[5];
|
||||
(void)result;
|
||||
} catch (...) {
|
||||
threw_exception = true;
|
||||
}
|
||||
// Note: Depending on implementation, this may or may not throw
|
||||
// We accept either behavior (return empty/invalid tensor or throw)
|
||||
(void)threw_exception; // Silence unused variable warning
|
||||
}
|
||||
|
||||
// ======================= Additional clamp edge case tests
|
||||
// =======================
|
||||
|
||||
TEST_F(TensorClampTest, ClampNoMinMax) {
|
||||
// Test clamp with no min and max (should be identity)
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor result = t.clamp(::std::optional<at::Scalar>(::std::nullopt),
|
||||
::std::optional<at::Scalar>(::std::nullopt));
|
||||
|
||||
ASSERT_EQ(result.numel(), 6);
|
||||
float* data = result.data_ptr<float>();
|
||||
for (int i = 0; i < 6; i++) {
|
||||
ASSERT_FLOAT_EQ(data[i], static_cast<float>(i));
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampOnlyMin) {
|
||||
// Test clamp with only min value
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor result =
|
||||
t.clamp(at::Scalar(2.5), ::std::optional<at::Scalar>(::std::nullopt));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.5f); // 0 < 2.5 -> 2.5
|
||||
ASSERT_FLOAT_EQ(data[1], 2.5f); // 1 < 2.5 -> 2.5
|
||||
ASSERT_FLOAT_EQ(data[2], 2.5f); // 2 < 2.5 -> 2.5
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampOnlyMax) {
|
||||
// Test clamp with only max value
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor result =
|
||||
t.clamp(::std::optional<at::Scalar>(::std::nullopt), at::Scalar(2.5));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[2], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[3], 2.5f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinOnlyTensor) {
|
||||
// Test clamp_min with Tensor
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor min_t = at::full({6}, 2.5f, at::kFloat);
|
||||
at::Tensor result = t.clamp_min(min_t);
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.5f); // 0 < 2.5 -> 2.5
|
||||
ASSERT_FLOAT_EQ(data[1], 2.5f); // 1 < 2.5 -> 2.5
|
||||
ASSERT_FLOAT_EQ(data[2], 2.5f); // 2 < 2.5 -> 2.5
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMaxOnlyTensor) {
|
||||
// Test clamp_max with Tensor
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor max_t = at::full({6}, 2.5f, at::kFloat);
|
||||
at::Tensor result = t.clamp_max(max_t);
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 1.0f);
|
||||
ASSERT_FLOAT_EQ(data[2], 2.0f);
|
||||
ASSERT_FLOAT_EQ(data[3], 2.5f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampWithTensorBothNone) {
|
||||
// Test clamp with both min and max as empty optional
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor result = t.clamp(::std::optional<at::Tensor>(::std::nullopt),
|
||||
::std::optional<at::Tensor>(::std::nullopt));
|
||||
|
||||
ASSERT_EQ(result.numel(), 6);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinTensorMaxNone) {
|
||||
// Test clamp with min tensor, max none
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor min_t = at::full({6}, 2.0f, at::kFloat);
|
||||
at::Tensor result = t.clamp(::std::optional<at::Tensor>(min_t),
|
||||
::std::optional<at::Tensor>(::std::nullopt));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampMinNoneMaxTensor) {
|
||||
// Test clamp with min none, max tensor
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
at::Tensor max_t = at::full({6}, 3.0f, at::kFloat);
|
||||
at::Tensor result = t.clamp(::std::optional<at::Tensor>(::std::nullopt),
|
||||
::std::optional<at::Tensor>(max_t));
|
||||
|
||||
float* data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[3], 3.0f);
|
||||
ASSERT_FLOAT_EQ(data[4], 3.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampInplaceMinNoneMax) {
|
||||
// Test clamp_ with min none
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
t.clamp_(::std::optional<at::Scalar>(::std::nullopt), at::Scalar(2.5));
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[3], 2.5f);
|
||||
}
|
||||
|
||||
TEST_F(TensorClampTest, ClampInplaceMaxNoneMin) {
|
||||
// Test clamp_ with max none
|
||||
at::Tensor t = at::arange(6, at::kFloat);
|
||||
t.clamp_(at::Scalar(2.0), ::std::optional<at::Scalar>(::std::nullopt));
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 2.0f);
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::Tensor::coalesce() and at::Tensor::is_coalesced()
|
||||
// ============================================================
|
||||
|
||||
// Helper: build a 2-D sparse COO tensor from indices and values.
|
||||
// indices shape: [sparse_dim, nnz], values shape: [nnz]
|
||||
static at::Tensor make_sparse(at::Tensor indices,
|
||||
at::Tensor values,
|
||||
c10::IntArrayRef size) {
|
||||
return at::sparse_coo_tensor(indices, values, size);
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, NewSparseNotCoalesced) {
|
||||
// A freshly created sparse COO tensor reports is_coalesced() == false.
|
||||
at::Tensor indices =
|
||||
at::tensor({0, 0, 1, 1, 1, 2}, at::kLong).reshape({2, 3});
|
||||
at::Tensor values = at::tensor({1.0f, 2.0f, 3.0f}, at::kFloat);
|
||||
at::Tensor sparse = make_sparse(indices, values, {3, 3});
|
||||
|
||||
ASSERT_FALSE(sparse.is_coalesced());
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, CoalesceReturnsSparse) {
|
||||
// coalesce() returns a sparse COO tensor.
|
||||
at::Tensor indices =
|
||||
at::tensor({0, 0, 1, 1, 1, 2}, at::kLong).reshape({2, 3});
|
||||
at::Tensor values = at::tensor({1.0f, 2.0f, 3.0f}, at::kFloat);
|
||||
at::Tensor sparse = make_sparse(indices, values, {3, 3});
|
||||
|
||||
at::Tensor coalesced = sparse.coalesce();
|
||||
|
||||
ASSERT_EQ(coalesced.layout(), c10::kSparse);
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, CoalescedTensorIsCoalesced) {
|
||||
// After calling coalesce(), is_coalesced() must return true.
|
||||
at::Tensor indices =
|
||||
at::tensor({0, 0, 1, 1, 1, 2}, at::kLong).reshape({2, 3});
|
||||
at::Tensor values = at::tensor({1.0f, 2.0f, 3.0f}, at::kFloat);
|
||||
at::Tensor sparse = make_sparse(indices, values, {3, 3});
|
||||
|
||||
at::Tensor coalesced = sparse.coalesce();
|
||||
|
||||
ASSERT_TRUE(coalesced.is_coalesced());
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, CoalesceDuplicateIndices_SumsValues) {
|
||||
// Duplicate indices [(0,1) appears twice] are merged; values are summed.
|
||||
// indices = [[0,0],[1,1]] (both at (0,1))
|
||||
at::Tensor indices = at::tensor({0, 0, 1, 1}, at::kLong).reshape({2, 2});
|
||||
at::Tensor values = at::tensor({1.0f, 2.0f}, at::kFloat);
|
||||
at::Tensor sparse = make_sparse(indices, values, {3, 3});
|
||||
|
||||
at::Tensor coalesced = sparse.coalesce();
|
||||
ASSERT_TRUE(coalesced.is_coalesced());
|
||||
// After coalescing, nnz should be 1 (duplicates merged)
|
||||
ASSERT_EQ(coalesced._nnz(), 1);
|
||||
// The merged value at (0,1) should be 1+2 = 3
|
||||
ASSERT_FLOAT_EQ(coalesced._values()[0].item<float>(), 3.0f);
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, CoalesceIdempotent) {
|
||||
// Calling coalesce() on an already-coalesced tensor returns the same tensor.
|
||||
at::Tensor indices = at::tensor({0, 1, 1, 2}, at::kLong).reshape({2, 2});
|
||||
at::Tensor values = at::tensor({1.0f, 2.0f}, at::kFloat);
|
||||
at::Tensor sparse = make_sparse(indices, values, {3, 3});
|
||||
|
||||
at::Tensor coalesced1 = sparse.coalesce();
|
||||
at::Tensor coalesced2 = coalesced1.coalesce(); // already coalesced
|
||||
|
||||
ASSERT_TRUE(coalesced2.is_coalesced());
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, CoalesceOnDenseTensorThrows) {
|
||||
// coalesce() on a dense tensor must throw.
|
||||
at::Tensor dense = at::ones({3, 3}, at::kFloat);
|
||||
ASSERT_THROW(dense.coalesce(), std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorCoalesceTest, IsCoalescedOnDenseTensorThrows) {
|
||||
// is_coalesced() on a dense tensor must throw.
|
||||
at::Tensor dense = at::ones({3, 3}, at::kFloat);
|
||||
ASSERT_THROW(dense.is_coalesced(), std::exception);
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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.
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
|
||||
#include <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/DeviceType.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::Tensor::cuda()
|
||||
// ============================================================
|
||||
|
||||
// After cuda(), the tensor should reside on a GPU device.
|
||||
TEST(TensorCudaTest, CpuTensorMovesToCuda) {
|
||||
at::Tensor cpu_t = at::tensor({1.0f, 2.0f, 3.0f}, at::kFloat);
|
||||
ASSERT_TRUE(cpu_t.is_cpu());
|
||||
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
ASSERT_TRUE(cuda_t.is_cuda());
|
||||
ASSERT_FALSE(cuda_t.is_cpu());
|
||||
}
|
||||
|
||||
// dtype and numel must be preserved.
|
||||
TEST(TensorCudaTest, DtypeAndNumelPreserved) {
|
||||
at::Tensor cpu_t = at::tensor({1, 2, 3, 4}, at::kInt);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
|
||||
ASSERT_EQ(cuda_t.scalar_type(), at::kInt);
|
||||
ASSERT_EQ(cuda_t.numel(), 4);
|
||||
}
|
||||
|
||||
// Values should round-trip back to CPU intact.
|
||||
TEST(TensorCudaTest, ValuesPreservedAfterRoundTrip) {
|
||||
std::vector<float> data = {1.0f, 2.5f, -3.0f, 4.75f};
|
||||
at::Tensor cpu_t = at::tensor(data, at::kFloat);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
at::Tensor back = cuda_t.cpu();
|
||||
|
||||
ASSERT_EQ(back.numel(), static_cast<int64_t>(data.size()));
|
||||
for (int64_t i = 0; i < back.numel(); ++i) {
|
||||
ASSERT_NEAR(back[i].item<float>(), data[static_cast<size_t>(i)], 1e-5f);
|
||||
}
|
||||
}
|
||||
|
||||
// shape (sizes) should be preserved.
|
||||
TEST(TensorCudaTest, ShapePreserved) {
|
||||
at::Tensor cpu_t = at::zeros({2, 3, 4}, at::kFloat);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
|
||||
ASSERT_EQ(cuda_t.dim(), 3);
|
||||
ASSERT_EQ(cuda_t.size(0), 2);
|
||||
ASSERT_EQ(cuda_t.size(1), 3);
|
||||
ASSERT_EQ(cuda_t.size(2), 4);
|
||||
}
|
||||
|
||||
// An already-CUDA tensor should still be CUDA after another cuda() call.
|
||||
TEST(TensorCudaTest, AlreadyCudaTensorStaysCuda) {
|
||||
at::Tensor cpu_t = at::tensor({7.0f}, at::kFloat);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
at::Tensor cuda_t2 = cuda_t.cuda();
|
||||
|
||||
ASSERT_TRUE(cuda_t2.is_cuda());
|
||||
ASSERT_NEAR(cuda_t2.cpu().item<float>(), 7.0f, 1e-6f);
|
||||
}
|
||||
|
||||
// device() should report a CUDA device.
|
||||
TEST(TensorCudaTest, DeviceIsCuda) {
|
||||
at::Tensor cpu_t = at::tensor({0.0f}, at::kFloat);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
|
||||
ASSERT_EQ(cuda_t.device().type(), c10::DeviceType::CUDA);
|
||||
}
|
||||
|
||||
TEST(TensorCudaTest, DefaultCudaUsesCurrentDevice) {
|
||||
if (c10::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
c10::cuda::CUDAGuard guard(1);
|
||||
at::Tensor cpu_t = at::tensor({1.0f}, at::kFloat);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
|
||||
ASSERT_EQ(cuda_t.device().type(), c10::DeviceType::CUDA);
|
||||
ASSERT_EQ(cuda_t.device().index(), 1);
|
||||
}
|
||||
|
||||
// is_cuda() / is_cpu() are mutually exclusive.
|
||||
TEST(TensorCudaTest, IsCudaAndIsCpuMutuallyExclusive) {
|
||||
at::Tensor cpu_t = at::tensor({1.0f, 2.0f}, at::kFloat);
|
||||
at::Tensor cuda_t = cpu_t.cuda();
|
||||
|
||||
ASSERT_TRUE(cuda_t.is_cuda());
|
||||
ASSERT_FALSE(cuda_t.is_cpu());
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,61 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
#include "utils/dense_sparse_conversion.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for compat::_PD_ConvertToSparseIfNeeded()
|
||||
// ============================================================
|
||||
|
||||
// Helper: create a small 2-D dense float tensor.
|
||||
static paddle::Tensor make_dense_2d(int rows = 3, int cols = 3) {
|
||||
// Use at::ones, but get the underlying paddle::Tensor.
|
||||
at::Tensor t = at::ones({rows, cols}, at::kFloat);
|
||||
return t._PD_GetInner();
|
||||
}
|
||||
|
||||
// ---- kStrided -> dense (no conversion) ----
|
||||
|
||||
TEST(DenseSparseConversionTest, kStrided_ReturnsDense) {
|
||||
paddle::Tensor dense = make_dense_2d();
|
||||
at::Tensor result = compat::_PD_ConvertToSparseIfNeeded(dense, c10::kStrided);
|
||||
|
||||
ASSERT_EQ(result.layout(), c10::kStrided);
|
||||
}
|
||||
|
||||
// ---- unsupported layout throws ----
|
||||
|
||||
TEST(DenseSparseConversionTest, UnsupportedLayout_Throws) {
|
||||
paddle::Tensor dense = make_dense_2d();
|
||||
|
||||
// kSparseBsr is not handled in the switch => PD_CHECK(false, ...) fires.
|
||||
ASSERT_THROW(compat::_PD_ConvertToSparseIfNeeded(dense, c10::kSparseBsr),
|
||||
std::exception);
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/ops/empty.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#ifdef PADDLE_WITH_XPU
|
||||
#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
|
||||
#endif
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ======================== at::empty basic tests ========================
|
||||
|
||||
TEST(ATenEmptyTest, BasicShape) {
|
||||
at::Tensor t = at::empty({3, 4});
|
||||
ASSERT_EQ(t.sizes()[0], 3);
|
||||
ASSERT_EQ(t.sizes()[1], 4);
|
||||
}
|
||||
|
||||
TEST(ATenEmptyTest, DtypeFloat) {
|
||||
at::Tensor t = at::empty({2, 2}, at::TensorOptions().dtype(at::kFloat));
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
}
|
||||
|
||||
TEST(ATenEmptyTest, DtypeDouble) {
|
||||
at::Tensor t = at::empty({4}, at::TensorOptions().dtype(at::kDouble));
|
||||
ASSERT_EQ(t.scalar_type(), at::kDouble);
|
||||
}
|
||||
|
||||
TEST(ATenEmptyTest, ExplicitArgsCpu) {
|
||||
// 6-argument overload: dtype, layout, device, pin_memory, memory_format
|
||||
at::Tensor t = at::empty(
|
||||
{2, 3}, at::kFloat, at::kStrided, at::kCPU, false, std::nullopt);
|
||||
ASSERT_EQ(t.sizes()[0], 2);
|
||||
ASSERT_EQ(t.sizes()[1], 3);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_FALSE(t.is_pinned());
|
||||
}
|
||||
|
||||
// ======================== pin_memory tests ========================
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
|
||||
// TensorOptions overload: pin_memory via options
|
||||
TEST(ATenEmptyTest, PinMemoryViaTensorOptions) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::TensorOptions opts =
|
||||
at::TensorOptions().dtype(at::kFloat).pinned_memory(true);
|
||||
at::Tensor t = at::empty({4, 4}, opts);
|
||||
ASSERT_TRUE(t.is_pinned())
|
||||
<< "Expected pinned memory tensor when TensorOptions.pinned_memory=true";
|
||||
}
|
||||
|
||||
// 6-argument overload: pin_memory = true (must use CPU device)
|
||||
TEST(ATenEmptyTest, PinMemoryViaExplicitArgs) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::Tensor t =
|
||||
at::empty({8}, at::kFloat, at::kStrided, at::kCPU, true, std::nullopt);
|
||||
ASSERT_TRUE(t.is_pinned())
|
||||
<< "Expected pinned memory tensor when pin_memory=true with CPU device";
|
||||
}
|
||||
|
||||
// pin_memory = false must NOT produce a pinned tensor
|
||||
TEST(ATenEmptyTest, NoPinMemoryViaExplicitArgs) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::Tensor t =
|
||||
at::empty({8}, at::kFloat, at::kStrided, at::kCUDA, false, std::nullopt);
|
||||
ASSERT_FALSE(t.is_pinned())
|
||||
<< "Expected non-pinned tensor when pin_memory=false";
|
||||
}
|
||||
|
||||
// Pinned tensor lives in pinned (host) memory, not on the GPU device itself
|
||||
TEST(ATenEmptyTest, PinnedTensorIsNotCuda) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::TensorOptions opts =
|
||||
at::TensorOptions().dtype(at::kFloat).pinned_memory(true);
|
||||
at::Tensor t = at::empty({16}, opts);
|
||||
ASSERT_TRUE(t.is_pinned());
|
||||
ASSERT_FALSE(t.is_cuda())
|
||||
<< "Pinned tensor should reside in host pinned memory, not on device";
|
||||
}
|
||||
|
||||
// Data pointer of a pinned tensor must be non-null
|
||||
TEST(ATenEmptyTest, PinnedTensorDataPtrNonNull) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::TensorOptions opts =
|
||||
at::TensorOptions().dtype(at::kFloat).pinned_memory(true);
|
||||
at::Tensor t = at::empty({32}, opts);
|
||||
ASSERT_TRUE(t.is_pinned());
|
||||
ASSERT_NE(t.data_ptr(), nullptr);
|
||||
}
|
||||
|
||||
TEST(ATenEmptyTest, DefaultCudaDeviceUsesCurrentDevice) {
|
||||
if (c10::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
c10::cuda::CUDAGuard guard(1);
|
||||
at::Tensor t =
|
||||
at::empty({8}, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
|
||||
ASSERT_TRUE(t.is_cuda());
|
||||
ASSERT_EQ(t.device().index(), 1);
|
||||
}
|
||||
|
||||
TEST(ATenEmptyTest, EmptyCudaHelperDefaultDeviceUsesCurrentDevice) {
|
||||
if (c10::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
c10::cuda::CUDAGuard guard(1);
|
||||
at::Tensor t = at::detail::empty_cuda(
|
||||
{8}, at::kFloat, at::Device(at::kCUDA), std::nullopt);
|
||||
|
||||
ASSERT_TRUE(t.is_cuda());
|
||||
ASSERT_EQ(t.device().index(), 1);
|
||||
}
|
||||
|
||||
TEST(ATenEmptyTest, EmptyCudaOptionsHelperDefaultDeviceUsesCurrentDevice) {
|
||||
if (c10::cuda::device_count() < 2) {
|
||||
return;
|
||||
}
|
||||
c10::cuda::CUDAGuard guard(1);
|
||||
at::Tensor t = at::detail::empty_cuda(
|
||||
{8}, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
|
||||
ASSERT_TRUE(t.is_cuda());
|
||||
ASSERT_EQ(t.device().index(), 1);
|
||||
}
|
||||
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
|
||||
#ifdef PADDLE_WITH_XPU
|
||||
TEST(ATenEmptyTest, DefaultXpuDeviceUsesCurrentDevice) {
|
||||
if (paddle::platform::GetXPUDeviceCount() < 2) {
|
||||
return;
|
||||
}
|
||||
paddle::platform::XPUDeviceGuard guard(1);
|
||||
at::Tensor t =
|
||||
at::empty({8}, at::TensorOptions().dtype(at::kFloat).device(at::kXPU));
|
||||
|
||||
ASSERT_EQ(t.device().type(), c10::DeviceType::XPU);
|
||||
ASSERT_EQ(t.device().index(), 1);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,53 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/ops/equal.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
TEST(TensorEqualTest, DifferentShapeReturnsFalse) {
|
||||
at::Tensor a = at::ones({2, 2}, at::kFloat);
|
||||
at::Tensor b = at::ones({2, 3}, at::kFloat);
|
||||
|
||||
ASSERT_FALSE(at::equal(a, b));
|
||||
ASSERT_FALSE(a.equal(b));
|
||||
}
|
||||
|
||||
TEST(TensorEqualTest, DtypeMismatchCastsOtherTensor) {
|
||||
at::Tensor a = at::tensor({1.0f, 2.0f, 3.0f}, at::kFloat);
|
||||
at::Tensor b = at::tensor({1, 2, 3}, at::kInt);
|
||||
|
||||
ASSERT_TRUE(at::equal(a, b));
|
||||
ASSERT_TRUE(a.equal(b));
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(TensorEqualTest, DeviceMismatchThrows) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::Tensor cpu = at::ones({2, 2}, at::kFloat);
|
||||
at::Tensor gpu =
|
||||
at::ones({2, 2}, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
|
||||
ASSERT_THROW((void)at::equal(cpu, gpu), std::exception);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,176 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "test/cpp/prim/init_env_utils.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
namespace {
|
||||
|
||||
class TensorExpandTest : public ::testing::Test {
|
||||
protected:
|
||||
static void SetUpTestSuite() { paddle::prim::InitTensorOperants(); }
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ======================== expand tests ========================
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandBasic) {
|
||||
// {3}.expand({3,4}) - PyTorch rejects non-singleton expansion (3 != 4)
|
||||
at::Tensor t = at::arange(3, at::kFloat);
|
||||
ASSERT_THROW(t.expand({3, 4}), std::exception);
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandSingleDim) {
|
||||
at::Tensor t = at::full({1}, 5.0f, at::kFloat);
|
||||
|
||||
at::Tensor result = t.expand({5});
|
||||
|
||||
ASSERT_EQ(result.numel(), 5);
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandMultipleDims) {
|
||||
// {1,3}.expand({2,3,4}) - PyTorch rejects non-singleton expansion (3 != 4)
|
||||
at::Tensor t = at::full({1, 3}, 1.0f, at::kFloat);
|
||||
ASSERT_THROW(t.expand({2, 3, 4}), std::exception);
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandWithImplicit) {
|
||||
// {3}.expand({3,4}) - PyTorch rejects non-singleton expansion (3 != 4)
|
||||
at::Tensor t = at::arange(3, at::kFloat);
|
||||
ASSERT_THROW(t.expand({3, 4}, true), std::exception);
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandPreservesValue) {
|
||||
// {3}.expand({3,4}) - PyTorch rejects non-singleton expansion (3 != 4)
|
||||
at::Tensor t = at::full({3}, 7.0f, at::kFloat);
|
||||
ASSERT_THROW(t.expand({3, 4}), std::exception);
|
||||
}
|
||||
|
||||
// ======================== expand_as tests ========================
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandAsBasic) {
|
||||
at::Tensor t = at::arange(3, at::kFloat).reshape({1, 3});
|
||||
at::Tensor other = at::zeros({2, 3}, at::kFloat);
|
||||
|
||||
at::Tensor result = t.expand_as(other);
|
||||
|
||||
ASSERT_EQ(result.sizes()[0], 2);
|
||||
ASSERT_EQ(result.sizes()[1], 3);
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandAsMatchSize) {
|
||||
at::Tensor t = at::full({1}, 7.0f, at::kFloat);
|
||||
at::Tensor other = at::zeros({3, 3, 3}, at::kFloat);
|
||||
|
||||
at::Tensor result = t.expand_as(other);
|
||||
|
||||
ASSERT_EQ(result.sizes().size(), 3);
|
||||
ASSERT_EQ(result.numel(), other.numel());
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandAsPreservesValue) {
|
||||
at::Tensor t = at::full({2, 1}, 5.0f, at::kFloat);
|
||||
at::Tensor other = at::zeros({2, 3}, at::kFloat);
|
||||
|
||||
at::Tensor result = t.expand_as(other);
|
||||
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 5.0f);
|
||||
}
|
||||
|
||||
// ======================== Additional tests for coverage
|
||||
// ========================
|
||||
|
||||
// Test tile fallback path when input_rank < target_rank
|
||||
// This triggers lines 86-100 in expand.h
|
||||
TEST_F(TensorExpandTest, ExpandTileFallbackLowRank) {
|
||||
// {2,1}.expand({1,4}) - PyTorch rejects shrinking non-singleton dims
|
||||
at::Tensor t = at::full({2, 1}, 1.0f, at::kFloat);
|
||||
ASSERT_THROW(t.expand({1, 4}), std::exception);
|
||||
}
|
||||
|
||||
// Test tile fallback when input_rank == target_rank
|
||||
// This triggers lines 119-130 in expand.h
|
||||
TEST_F(TensorExpandTest, ExpandSameRankTileFallback) {
|
||||
// {2,3}.expand({2,6}) - PyTorch only allows expanding singleton dims
|
||||
at::Tensor t = at::full({2, 3}, 2.0f, at::kFloat);
|
||||
ASSERT_THROW(t.expand({2, 6}), std::exception);
|
||||
}
|
||||
|
||||
// Test zero dimension handling
|
||||
// This triggers lines 90-94 and 122-126 in expand.h
|
||||
TEST_F(TensorExpandTest, ExpandZeroDim) {
|
||||
// {0}.expand({0,3}) - PyTorch rejects non-singleton expansion (0 != 3)
|
||||
at::Tensor t = at::full({0}, 1.0f, at::kFloat);
|
||||
ASSERT_THROW(t.expand({0, 3}), std::exception);
|
||||
}
|
||||
|
||||
// Test input_rank > target_rank branch
|
||||
// This triggers lines 131-136 in expand.h
|
||||
TEST_F(TensorExpandTest, ExpandHighRankToLowRank) {
|
||||
// Input has more dimensions than target - PyTorch rejects this
|
||||
at::Tensor t = at::full({2, 3, 4}, 1.0f, at::kFloat);
|
||||
ASSERT_THROW(t.expand({3, 4}), std::exception);
|
||||
}
|
||||
|
||||
// Test expand_as with tile fallback
|
||||
TEST_F(TensorExpandTest, ExpandAsTileFallback) {
|
||||
// {2,1}.expand_as({1,4}) - PyTorch rejects shrinking non-singleton dims
|
||||
at::Tensor t = at::full({2, 1}, 3.0f, at::kFloat);
|
||||
at::Tensor other = at::zeros({1, 4}, at::kFloat);
|
||||
|
||||
ASSERT_THROW(t.expand_as(other), std::exception);
|
||||
}
|
||||
|
||||
// Test preserve non-singleton dimension (matching dimension)
|
||||
TEST_F(TensorExpandTest, ExpandPreserveNonSingleton) {
|
||||
// {3,1}.expand({3,4}) - dim 0 matches (3), dim 1 expands (1->4)
|
||||
at::Tensor t = at::full({3, 1}, 5.0f, at::kFloat);
|
||||
at::Tensor result = t.expand({3, 4});
|
||||
|
||||
ASSERT_EQ(result.sizes()[0], 3);
|
||||
ASSERT_EQ(result.sizes()[1], 4);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 5.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[3], 5.0f);
|
||||
}
|
||||
|
||||
// Test expand function (not member function)
|
||||
TEST_F(TensorExpandTest, ExpandFunction) {
|
||||
at::Tensor t = at::full({1}, 7.0f, at::kFloat);
|
||||
|
||||
at::Tensor result = at::expand(t, {3, 4});
|
||||
|
||||
ASSERT_EQ(result.sizes()[0], 3);
|
||||
ASSERT_EQ(result.sizes()[1], 4);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 7.0f);
|
||||
}
|
||||
|
||||
TEST_F(TensorExpandTest, ExpandAsMemberFunction) {
|
||||
at::Tensor t = at::full({1, 2}, 4.0f, at::kFloat);
|
||||
at::Tensor other = at::zeros({3, 2}, at::kFloat);
|
||||
|
||||
at::Tensor result = t.expand_as(other);
|
||||
|
||||
ASSERT_EQ(result.sizes()[0], 3);
|
||||
ASSERT_EQ(result.sizes()[1], 2);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 4.0f);
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::eye()
|
||||
// ============================================================
|
||||
|
||||
// Helper: verify that a 2-D tensor is an identity-like matrix
|
||||
// (diagonal == 1, off-diagonal == 0).
|
||||
static void CheckEye(const at::Tensor& t, int64_t rows, int64_t cols) {
|
||||
ASSERT_EQ(t.dim(), 2);
|
||||
ASSERT_EQ(t.size(0), rows);
|
||||
ASSERT_EQ(t.size(1), cols);
|
||||
for (int64_t i = 0; i < rows; ++i) {
|
||||
for (int64_t j = 0; j < cols; ++j) {
|
||||
float expected = (i == j) ? 1.0f : 0.0f;
|
||||
ASSERT_FLOAT_EQ(t[i][j].item<float>(), expected)
|
||||
<< "Mismatch at (" << i << ", " << j << ")";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- eye(n) -------------------------------------------------------
|
||||
|
||||
TEST(ATenEyeTest, SquareDefaultDtype) {
|
||||
// eye(n) should produce an n×n float32 identity matrix.
|
||||
at::Tensor t = at::eye(4);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 4, 4);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, SquareTensorOptionsFloat) {
|
||||
// eye(n, TensorOptions) — explicit float32.
|
||||
at::Tensor t = at::eye(3, at::TensorOptions().dtype(at::kFloat));
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 3, 3);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, SquareTensorOptionsDouble) {
|
||||
// eye(n, TensorOptions) — explicit float64.
|
||||
at::Tensor t = at::eye(5, at::TensorOptions().dtype(at::kDouble));
|
||||
ASSERT_EQ(t.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(t.size(0), 5);
|
||||
ASSERT_EQ(t.size(1), 5);
|
||||
for (int64_t i = 0; i < 5; ++i) {
|
||||
ASSERT_DOUBLE_EQ(t[i][i].item<double>(), 1.0);
|
||||
if (i + 1 < 5) {
|
||||
ASSERT_DOUBLE_EQ(t[i][i + 1].item<double>(), 0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// eye(n, dtype, layout, device, pin_memory) — separate-params overload
|
||||
|
||||
TEST(ATenEyeTest, SquareSeparateParamsFloat) {
|
||||
at::Tensor t =
|
||||
at::eye(4, at::kFloat, /*layout=*/std::nullopt, at::kCPU, false);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 4, 4);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, SquareSeparateParamsNulloptDtype) {
|
||||
// When dtype is nullopt the default dtype (float32) should be used.
|
||||
at::Tensor t =
|
||||
at::eye(3, std::nullopt, /*layout=*/std::nullopt, at::kCPU, false);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 3, 3);
|
||||
}
|
||||
|
||||
// ---- eye(n, m) -------------------------------------------------------
|
||||
|
||||
TEST(ATenEyeTest, RectangularWiderThanTall) {
|
||||
// n < m: identity portion fits entirely within row range.
|
||||
at::Tensor t = at::eye(3, 5);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 3, 5);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, RectangularTallerThanWide) {
|
||||
// n > m: identity portion fits entirely within column range.
|
||||
at::Tensor t = at::eye(5, 3);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 5, 3);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, RectangularSquareEquivalent) {
|
||||
// eye(n, n) should behave like eye(n).
|
||||
at::Tensor t2 = at::eye(4, 4);
|
||||
at::Tensor t1 = at::eye(4);
|
||||
CheckEye(t2, 4, 4);
|
||||
for (int64_t i = 0; i < 4; ++i)
|
||||
for (int64_t j = 0; j < 4; ++j)
|
||||
ASSERT_FLOAT_EQ(t1[i][j].item<float>(), t2[i][j].item<float>());
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, RectangularTensorOptionsDouble) {
|
||||
// eye(n, m, TensorOptions) — float64.
|
||||
at::Tensor t = at::eye(2, 4, at::TensorOptions().dtype(at::kDouble));
|
||||
ASSERT_EQ(t.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(t.size(0), 2);
|
||||
ASSERT_EQ(t.size(1), 4);
|
||||
ASSERT_DOUBLE_EQ(t[0][0].item<double>(), 1.0);
|
||||
ASSERT_DOUBLE_EQ(t[1][1].item<double>(), 1.0);
|
||||
ASSERT_DOUBLE_EQ(t[0][1].item<double>(), 0.0);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, RectangularSeparateParams) {
|
||||
// eye(n, m, dtype, layout, device, pin_memory)
|
||||
at::Tensor t =
|
||||
at::eye(3, 5, at::kDouble, /*layout=*/std::nullopt, at::kCPU, false);
|
||||
ASSERT_EQ(t.scalar_type(), at::kDouble);
|
||||
CheckEye(t, 3, 5);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, RectangularSeparateParamsNulloptDtype) {
|
||||
at::Tensor t =
|
||||
at::eye(4, 6, std::nullopt, /*layout=*/std::nullopt, at::kCPU, false);
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
CheckEye(t, 4, 6);
|
||||
}
|
||||
|
||||
// ---- 1×1 edge case -------------------------------------------------------
|
||||
|
||||
TEST(ATenEyeTest, OneByOne) {
|
||||
at::Tensor t = at::eye(1);
|
||||
ASSERT_EQ(t.numel(), 1);
|
||||
ASSERT_FLOAT_EQ(t[0][0].item<float>(), 1.0f);
|
||||
}
|
||||
|
||||
// ---- GPU tests (compiled only when CUDA / HIP is available) --------------
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(ATenEyeTest, SquareOnGPU) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::Tensor t =
|
||||
at::eye(4, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
at::Tensor t_cpu = t.to(at::kCPU);
|
||||
CheckEye(t_cpu, 4, 4);
|
||||
}
|
||||
|
||||
TEST(ATenEyeTest, RectangularOnGPU) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
at::Tensor t =
|
||||
at::eye(3, 5, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA));
|
||||
at::Tensor t_cpu = t.to(at::kCPU);
|
||||
CheckEye(t_cpu, 3, 5);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,174 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/ops/arange.h>
|
||||
#include <ATen/ops/empty.h>
|
||||
#include <ATen/ops/full.h>
|
||||
#include <ATen/ops/ones.h>
|
||||
#include <ATen/ops/zeros.h>
|
||||
#include <c10/core/DefaultDtype.h>
|
||||
#include <c10/core/ScalarTypeToTypeMeta.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace {
|
||||
|
||||
class DefaultDtypeGuard {
|
||||
public:
|
||||
explicit DefaultDtypeGuard(c10::ScalarType dtype)
|
||||
: previous_(c10::get_default_dtype()) {
|
||||
c10::set_default_dtype(c10::scalarTypeToTypeMeta(dtype));
|
||||
}
|
||||
|
||||
~DefaultDtypeGuard() { c10::set_default_dtype(previous_); }
|
||||
|
||||
private:
|
||||
caffe2::TypeMeta previous_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest, EmptyNulloptDtypeUsesCurrentDefault) {
|
||||
DefaultDtypeGuard guard(at::kDouble);
|
||||
|
||||
at::Tensor tensor = at::empty(
|
||||
{2, 3}, std::nullopt, at::kStrided, at::kCPU, false, std::nullopt);
|
||||
|
||||
ASSERT_EQ(tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(tensor.sizes(), c10::IntArrayRef({2, 3}));
|
||||
}
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest, ArangeOmittedDtypeUsesLongForIntegralInputs) {
|
||||
DefaultDtypeGuard guard(at::kDouble);
|
||||
|
||||
at::Tensor end_only_default = at::arange(5);
|
||||
at::Tensor start_end_default = at::arange(1, 6);
|
||||
at::Tensor start_end_step_default = at::arange(1, 7, 2);
|
||||
at::Tensor end_only_nullopt =
|
||||
at::arange(5, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
at::Tensor start_end_nullopt =
|
||||
at::arange(1, 6, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
at::Tensor start_end_step_nullopt =
|
||||
at::arange(1, 7, 2, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
|
||||
ASSERT_EQ(end_only_default.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(start_end_default.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(start_end_step_default.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(end_only_nullopt.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(start_end_nullopt.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(start_end_step_nullopt.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(end_only_default.data_ptr<int64_t>()[4], 4);
|
||||
ASSERT_EQ(start_end_default.data_ptr<int64_t>()[0], 1);
|
||||
ASSERT_EQ(start_end_step_default.data_ptr<int64_t>()[2], 5);
|
||||
ASSERT_EQ(end_only_nullopt.data_ptr<int64_t>()[4], 4);
|
||||
ASSERT_EQ(start_end_nullopt.data_ptr<int64_t>()[0], 1);
|
||||
ASSERT_EQ(start_end_step_nullopt.data_ptr<int64_t>()[2], 5);
|
||||
}
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest,
|
||||
ArangeOmittedDtypeKeepsLargeInt64InputsExact) {
|
||||
constexpr int64_t kStart = (1LL << 53) + 1;
|
||||
constexpr int64_t kEnd = kStart + 4;
|
||||
|
||||
at::Tensor by_default = at::arange(kStart, kEnd);
|
||||
at::Tensor by_nullopt =
|
||||
at::arange(kStart, kEnd, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
|
||||
ASSERT_EQ(by_default.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(by_nullopt.scalar_type(), at::kLong);
|
||||
ASSERT_EQ(by_default.numel(), 4);
|
||||
ASSERT_EQ(by_nullopt.numel(), 4);
|
||||
|
||||
for (int64_t i = 0; i < 4; ++i) {
|
||||
ASSERT_EQ(by_default.data_ptr<int64_t>()[i], kStart + i);
|
||||
ASSERT_EQ(by_nullopt.data_ptr<int64_t>()[i], kStart + i);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest,
|
||||
ArangeOmittedDtypeUsesCurrentDefaultForFloatingInputs) {
|
||||
DefaultDtypeGuard guard(at::kDouble);
|
||||
|
||||
at::Tensor end_only_default = at::arange(5.0);
|
||||
at::Tensor start_end_default = at::arange(1.0, 6.0);
|
||||
at::Tensor start_end_step_default = at::arange(1.0, 7.0, 2.0);
|
||||
at::Tensor end_only_nullopt =
|
||||
at::arange(5.0, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
at::Tensor start_end_nullopt =
|
||||
at::arange(1.0, 6.0, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
at::Tensor start_end_step_nullopt =
|
||||
at::arange(1.0, 7.0, 2.0, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
|
||||
ASSERT_EQ(end_only_default.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(start_end_default.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(start_end_step_default.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(end_only_nullopt.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(start_end_nullopt.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(start_end_step_nullopt.scalar_type(), at::kDouble);
|
||||
ASSERT_DOUBLE_EQ(end_only_default.data_ptr<double>()[4], 4.0);
|
||||
ASSERT_DOUBLE_EQ(start_end_default.data_ptr<double>()[0], 1.0);
|
||||
ASSERT_DOUBLE_EQ(start_end_step_default.data_ptr<double>()[2], 5.0);
|
||||
ASSERT_DOUBLE_EQ(end_only_nullopt.data_ptr<double>()[4], 4.0);
|
||||
ASSERT_DOUBLE_EQ(start_end_nullopt.data_ptr<double>()[0], 1.0);
|
||||
ASSERT_DOUBLE_EQ(start_end_step_nullopt.data_ptr<double>()[2], 5.0);
|
||||
}
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest, FullNulloptDtypeUsesCurrentDefault) {
|
||||
DefaultDtypeGuard guard(at::kDouble);
|
||||
|
||||
at::Tensor tensor =
|
||||
at::full({2, 3}, 1.25, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
at::Tensor symint_tensor = at::full_symint(c10::SymIntArrayRef({2, 3}),
|
||||
2.5,
|
||||
std::nullopt,
|
||||
std::nullopt,
|
||||
at::kCPU,
|
||||
false);
|
||||
|
||||
ASSERT_EQ(tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(symint_tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_DOUBLE_EQ(tensor.data_ptr<double>()[0], 1.25);
|
||||
ASSERT_DOUBLE_EQ(symint_tensor.data_ptr<double>()[0], 2.5);
|
||||
}
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest, OnesNulloptDtypeUsesCurrentDefault) {
|
||||
DefaultDtypeGuard guard(at::kDouble);
|
||||
|
||||
at::Tensor tensor =
|
||||
at::ones({2, 3}, std::nullopt, std::nullopt, at::kCPU, false);
|
||||
at::Tensor symint_tensor = at::ones_symint(
|
||||
c10::SymIntArrayRef({2, 3}), std::nullopt, std::nullopt, at::kCPU, false);
|
||||
|
||||
ASSERT_EQ(tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(symint_tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_DOUBLE_EQ(tensor.data_ptr<double>()[0], 1.0);
|
||||
ASSERT_DOUBLE_EQ(symint_tensor.data_ptr<double>()[0], 1.0);
|
||||
}
|
||||
|
||||
TEST(ATenFactoryDefaultDtypeTest, ZerosNulloptDtypeUsesCurrentDefault) {
|
||||
DefaultDtypeGuard guard(at::kDouble);
|
||||
|
||||
at::Tensor tensor =
|
||||
at::zeros({2, 3}, std::nullopt, at::kStrided, at::kCPU, false);
|
||||
at::Tensor symint_tensor = at::zeros_symint(
|
||||
c10::SymIntArrayRef({2, 3}), std::nullopt, at::kStrided, at::kCPU, false);
|
||||
|
||||
ASSERT_EQ(tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_EQ(symint_tensor.scalar_type(), at::kDouble);
|
||||
ASSERT_DOUBLE_EQ(tensor.data_ptr<double>()[0], 0.0);
|
||||
ASSERT_DOUBLE_EQ(symint_tensor.data_ptr<double>()[0], 0.0);
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================================
|
||||
// Flatten Tests
|
||||
// ============================================================================
|
||||
|
||||
TEST(TestFlatten, FlattenAllDims) {
|
||||
// Test flatten with start_dim=0, end_dim=-1
|
||||
// Flattens the entire tensor to 1D
|
||||
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
||||
at::Tensor flattened = tensor.flatten(0, -1);
|
||||
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({24}));
|
||||
ASSERT_EQ(flattened.numel(), tensor.numel());
|
||||
}
|
||||
|
||||
TEST(TestFlatten, FlattenPartialDims) {
|
||||
// Test flatten with specific start and end dimensions
|
||||
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
||||
|
||||
// Flatten dimensions 1 to 2 (3*4 = 12)
|
||||
at::Tensor flattened = tensor.flatten(1, 2);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12, 5}));
|
||||
ASSERT_EQ(flattened.numel(), tensor.numel());
|
||||
}
|
||||
|
||||
TEST(TestFlatten, FlattenSingleDim) {
|
||||
// Test flatten when start_dim == end_dim (should be no-op)
|
||||
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
||||
|
||||
at::Tensor flattened = tensor.flatten(1, 1);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
}
|
||||
|
||||
TEST(TestFlatten, FlattenNegativeDims) {
|
||||
// Test flatten with negative dimension indices
|
||||
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
||||
|
||||
// Flatten from -3 to -2 (dimensions 1 to 2)
|
||||
at::Tensor flattened = tensor.flatten(-3, -2);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12, 5}));
|
||||
}
|
||||
|
||||
TEST(TestFlatten, FlattenFirstTwoDims) {
|
||||
// Test flatten on first two dimensions
|
||||
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
||||
|
||||
at::Tensor flattened = tensor.flatten(0, 1);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({6, 4}));
|
||||
}
|
||||
|
||||
TEST(TestFlatten, FlattenLastTwoDims) {
|
||||
// Test flatten on last two dimensions
|
||||
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
||||
|
||||
at::Tensor flattened = tensor.flatten(1, 2);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12}));
|
||||
}
|
||||
|
||||
TEST(TestFlatten, FlattenDataIntegrity) {
|
||||
// Test that flatten preserves data
|
||||
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
||||
at::Tensor flattened = tensor.flatten(0, -1);
|
||||
|
||||
const float* original_data = tensor.data_ptr<float>();
|
||||
const float* flattened_data = flattened.data_ptr<float>();
|
||||
|
||||
for (int64_t i = 0; i < tensor.numel(); ++i) {
|
||||
ASSERT_EQ(original_data[i], flattened_data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Unflatten Tests
|
||||
// ============================================================================
|
||||
|
||||
TEST(TestUnflatten, UnflattenBasic) {
|
||||
// Test basic unflatten operation
|
||||
at::Tensor tensor = at::ones({4, 6, 8}, at::kFloat);
|
||||
|
||||
// Unflatten dimension 1 (size 6) into (2, 3)
|
||||
at::Tensor unflattened = tensor.unflatten(1, c10::IntArrayRef({2, 3}));
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({4, 2, 3, 8}));
|
||||
ASSERT_EQ(unflattened.numel(), tensor.numel());
|
||||
}
|
||||
|
||||
TEST(TestUnflatten, UnflattenFirstDim) {
|
||||
// Test unflatten on first dimension
|
||||
at::Tensor tensor = at::ones({6, 4}, at::kFloat);
|
||||
|
||||
// Unflatten dimension 0 (size 6) into (2, 3)
|
||||
at::Tensor unflattened = tensor.unflatten(0, c10::IntArrayRef({2, 3}));
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
}
|
||||
|
||||
TEST(TestUnflatten, UnflattenLastDim) {
|
||||
// Test unflatten on last dimension
|
||||
at::Tensor tensor = at::ones({2, 12}, at::kFloat);
|
||||
|
||||
// Unflatten dimension 1 (size 12) into (3, 4)
|
||||
at::Tensor unflattened = tensor.unflatten(1, c10::IntArrayRef({3, 4}));
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
}
|
||||
|
||||
TEST(TestUnflatten, UnflattenNegativeDim) {
|
||||
// Test unflatten with negative dimension index
|
||||
at::Tensor tensor = at::ones({4, 6, 8}, at::kFloat);
|
||||
|
||||
// Unflatten dimension -1 (last dim, size 8) into (4, 2)
|
||||
at::Tensor unflattened = tensor.unflatten(-1, c10::IntArrayRef({4, 2}));
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({4, 6, 4, 2}));
|
||||
}
|
||||
|
||||
TEST(TestUnflatten, UnflattenSymInt) {
|
||||
// Test unflatten_symint (should behave same as unflatten)
|
||||
at::Tensor tensor = at::ones({4, 6, 8}, at::kFloat);
|
||||
|
||||
// Unflatten dimension 1 using symint version
|
||||
// Note: Must keep the underlying data alive
|
||||
std::vector<c10::SymInt> sizes_vec = {2, 3};
|
||||
c10::SymIntArrayRef sizes(sizes_vec);
|
||||
at::Tensor unflattened = tensor.unflatten_symint(1, sizes);
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({4, 2, 3, 8}));
|
||||
}
|
||||
|
||||
TEST(TestUnflatten, UnflattenDataIntegrity) {
|
||||
// Test that unflatten preserves data
|
||||
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 12});
|
||||
at::Tensor unflattened = tensor.unflatten(1, c10::IntArrayRef({3, 4}));
|
||||
|
||||
// Verify shape
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
|
||||
// Verify numel
|
||||
ASSERT_EQ(unflattened.numel(), tensor.numel());
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Flatten and Unflatten Combined Tests
|
||||
// ============================================================================
|
||||
|
||||
TEST(TestFlattenUnflatten, RoundTrip) {
|
||||
// Test that flatten followed by unflatten restores original shape
|
||||
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
||||
|
||||
// Flatten dimensions 1 and 2
|
||||
at::Tensor flattened = tensor.flatten(1, 2);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12}));
|
||||
|
||||
// Unflatten back to original shape
|
||||
at::Tensor unflattened = flattened.unflatten(1, c10::IntArrayRef({3, 4}));
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
||||
|
||||
// Verify data integrity
|
||||
ASSERT_EQ(tensor.numel(), unflattened.numel());
|
||||
}
|
||||
|
||||
TEST(TestFlattenUnflatten, MultipleOperations) {
|
||||
// Test multiple flatten/unflatten operations
|
||||
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
||||
|
||||
// Flatten all dimensions
|
||||
at::Tensor flattened = tensor.flatten(0, -1);
|
||||
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({120}));
|
||||
|
||||
// Unflatten into different shape
|
||||
at::Tensor unflattened = flattened.unflatten(0, c10::IntArrayRef({6, 20}));
|
||||
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({6, 20}));
|
||||
|
||||
// Unflatten again
|
||||
at::Tensor final_tensor = unflattened.unflatten(1, c10::IntArrayRef({4, 5}));
|
||||
ASSERT_EQ(final_tensor.sizes(), c10::IntArrayRef({6, 4, 5}));
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/ops/from_blob.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/common/macros.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_stride_kernel);
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
#include <cuda_runtime.h>
|
||||
#elif defined(PADDLE_WITH_HIP)
|
||||
#include <hip/hip_runtime.h>
|
||||
#endif
|
||||
|
||||
// ======================== CPU place detection ========================
|
||||
|
||||
// No device specified: CPU pointer → tensor must be on CPU.
|
||||
TEST(ATenFromBlobTest, CpuPtrDefaultsToCpu) {
|
||||
float data[4] = {1.0f, 2.0f, 3.0f, 4.0f};
|
||||
at::Tensor t = at::from_blob(data, {4}, at::kFloat);
|
||||
ASSERT_TRUE(t.is_cpu());
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_EQ(t.numel(), 4);
|
||||
}
|
||||
|
||||
// Explicitly pass CPU options: still CPU.
|
||||
TEST(ATenFromBlobTest, CpuPtrWithCpuOptions) {
|
||||
float data[3] = {1.0f, 2.0f, 3.0f};
|
||||
at::Tensor t = at::from_blob(
|
||||
data, {3}, at::TensorOptions().dtype(at::kFloat).device(at::kCPU));
|
||||
ASSERT_TRUE(t.is_cpu());
|
||||
}
|
||||
|
||||
// Data pointer must be preserved (no copy).
|
||||
TEST(ATenFromBlobTest, DataPtrPreserved) {
|
||||
float data[4] = {10.f, 20.f, 30.f, 40.f};
|
||||
at::Tensor t = at::from_blob(data, {4}, at::kFloat);
|
||||
ASSERT_EQ(t.data_ptr<float>(), data);
|
||||
}
|
||||
|
||||
// Shape and strides are correctly set.
|
||||
TEST(ATenFromBlobTest, ShapeAndStrides) {
|
||||
float data[6] = {};
|
||||
at::Tensor t = at::from_blob(data, {2, 3}, at::kFloat);
|
||||
ASSERT_EQ(t.sizes()[0], 2);
|
||||
ASSERT_EQ(t.sizes()[1], 3);
|
||||
// contiguous strides: [3, 1]
|
||||
ASSERT_EQ(t.strides()[0], 3);
|
||||
ASSERT_EQ(t.strides()[1], 1);
|
||||
}
|
||||
|
||||
// Explicit strides overload.
|
||||
TEST(ATenFromBlobTest, ExplicitStrides) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
// Row-major 2×3 laid out in memory, but we interpret as column-major strides
|
||||
float data[6] = {1, 2, 3, 4, 5, 6};
|
||||
at::Tensor t = at::from_blob(data, {2, 3}, {1, 2}, at::kFloat);
|
||||
ASSERT_EQ(t.strides()[0], 1);
|
||||
ASSERT_EQ(t.strides()[1], 2);
|
||||
ASSERT_TRUE(t.is_cpu());
|
||||
}
|
||||
|
||||
// Deleter is called when the tensor is destroyed.
|
||||
TEST(ATenFromBlobTest, DeleterCalled) {
|
||||
bool deleted = false;
|
||||
{
|
||||
float* data = new float[4]{};
|
||||
at::Tensor t = at::from_blob(
|
||||
data,
|
||||
{4},
|
||||
[&deleted](void* p) {
|
||||
deleted = true;
|
||||
delete[] static_cast<float*>(p);
|
||||
},
|
||||
at::kFloat);
|
||||
ASSERT_FALSE(deleted);
|
||||
}
|
||||
ASSERT_TRUE(deleted);
|
||||
}
|
||||
|
||||
// Deleter + strides overload.
|
||||
TEST(ATenFromBlobTest, DeleterWithStrides) {
|
||||
bool deleted = false;
|
||||
{
|
||||
float* data = new float[6]{};
|
||||
at::Tensor t = at::from_blob(
|
||||
data,
|
||||
{2, 3},
|
||||
{3, 1},
|
||||
[&deleted](void* p) {
|
||||
deleted = true;
|
||||
delete[] static_cast<float*>(p);
|
||||
},
|
||||
at::kFloat);
|
||||
ASSERT_FALSE(deleted);
|
||||
ASSERT_TRUE(t.is_cpu());
|
||||
}
|
||||
ASSERT_TRUE(deleted);
|
||||
}
|
||||
|
||||
// ======================== GPU place detection ========================
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
|
||||
// No device specified: GPU pointer → tensor must be on CUDA automatically.
|
||||
TEST(ATenFromBlobTest, GpuPtrDefaultsToCuda) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
float* d_data = nullptr;
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
cudaMalloc(&d_data, 4 * sizeof(float));
|
||||
#else
|
||||
hipMalloc(&d_data, 4 * sizeof(float));
|
||||
#endif
|
||||
|
||||
at::Tensor t = at::from_blob(d_data, {4}, at::kFloat);
|
||||
ASSERT_TRUE(t.is_cuda())
|
||||
<< "Expected GPU tensor when data pointer lives on device";
|
||||
ASSERT_FALSE(t.is_cpu());
|
||||
ASSERT_EQ(t.scalar_type(), at::kFloat);
|
||||
ASSERT_EQ(t.numel(), 4);
|
||||
ASSERT_EQ(t.data_ptr<float>(), d_data);
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
cudaFree(d_data);
|
||||
#else
|
||||
hipFree(d_data);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Explicit CUDA device option + GPU pointer → still CUDA.
|
||||
TEST(ATenFromBlobTest, GpuPtrWithCudaOptions) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
float* d_data = nullptr;
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
cudaMalloc(&d_data, 4 * sizeof(float));
|
||||
#else
|
||||
hipMalloc(&d_data, 4 * sizeof(float));
|
||||
#endif
|
||||
|
||||
at::Tensor t = at::from_blob(
|
||||
d_data, {4}, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0));
|
||||
ASSERT_TRUE(t.is_cuda());
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
cudaFree(d_data);
|
||||
#else
|
||||
hipFree(d_data);
|
||||
#endif
|
||||
}
|
||||
|
||||
// target_device overrides auto-detection.
|
||||
TEST(ATenFromBlobTest, TargetDeviceOverride) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
float* d_data = nullptr;
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
cudaMalloc(&d_data, 4 * sizeof(float));
|
||||
#else
|
||||
hipMalloc(&d_data, 4 * sizeof(float));
|
||||
#endif
|
||||
|
||||
at::Tensor t = at::for_blob(d_data, {4})
|
||||
.options(at::kFloat)
|
||||
.target_device(at::Device(at::kCUDA, 0))
|
||||
.make_tensor();
|
||||
ASSERT_TRUE(t.is_cuda());
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
cudaFree(d_data);
|
||||
#else
|
||||
hipFree(d_data);
|
||||
#endif
|
||||
}
|
||||
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
@@ -0,0 +1,50 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ======================== register_hook tests ========================
|
||||
|
||||
TEST(TensorHookTest, RegisterHookThrows) {
|
||||
// register_hook should throw exception as Paddle doesn't support hooks
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
auto hook = [](const at::Tensor& grad) { return grad; };
|
||||
EXPECT_THROW(t.register_hook(hook), std::runtime_error);
|
||||
}
|
||||
|
||||
TEST(TensorHookTest, RegisterHookWithLambda) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
// Lambda that captures nothing
|
||||
EXPECT_THROW(t.register_hook([](const at::Tensor&) { return at::Tensor(); }),
|
||||
std::runtime_error);
|
||||
}
|
||||
|
||||
TEST(TensorHookTest, RegisterHookWithMoveOnly) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
|
||||
// Move-only lambda
|
||||
EXPECT_THROW(
|
||||
t.register_hook([](const at::Tensor&) mutable { return at::Tensor(); }),
|
||||
std::runtime_error);
|
||||
}
|
||||
@@ -0,0 +1,447 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/TensorIndexing.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/List.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/common/macros.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_stride_kernel);
|
||||
|
||||
// ======================== index tests ========================
|
||||
|
||||
TEST(TensorIndexTest, IndexWithSingleTensor) {
|
||||
// Create tensor [0, 10, 20, 30, 40]
|
||||
at::Tensor t = at::arange(5, at::kFloat);
|
||||
for (int i = 0; i < 5; i++) {
|
||||
t.data_ptr<float>()[i] = static_cast<float>(i * 10);
|
||||
}
|
||||
|
||||
// Index with [0, 2, 4]
|
||||
at::Tensor idx = at::empty({3}, at::kLong);
|
||||
int64_t* idx_data = idx.data_ptr<int64_t>();
|
||||
idx_data[0] = 0;
|
||||
idx_data[1] = 2;
|
||||
idx_data[2] = 4;
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx);
|
||||
|
||||
at::Tensor result = t.index(indices);
|
||||
ASSERT_EQ(result.numel(), 3);
|
||||
|
||||
float* result_data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(result_data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(result_data[1], 20.0f);
|
||||
ASSERT_FLOAT_EQ(result_data[2], 40.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, SliceKeepsStrideWithoutContiguousCopy) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
at::Tensor base = at::arange(24, at::kFloat).reshape({4, 6});
|
||||
at::Tensor transposed = base.t(); // shape: [6, 4], strides: [1, 6]
|
||||
ASSERT_FALSE(transposed.is_contiguous());
|
||||
|
||||
at::Tensor sliced =
|
||||
transposed.index({at::indexing::Slice(1, 5), at::indexing::Slice(0, 3)});
|
||||
|
||||
ASSERT_EQ(sliced.sizes(), c10::IntArrayRef({4, 3}));
|
||||
ASSERT_EQ(sliced.strides(), c10::IntArrayRef({1, 6}));
|
||||
ASSERT_EQ(sliced.stride(0), transposed.stride(0));
|
||||
ASSERT_EQ(sliced.stride(1), transposed.stride(1));
|
||||
ASSERT_FALSE(sliced.is_contiguous());
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, IndexWithEmptyInitializerListReturnsSelf) {
|
||||
at::Tensor t = at::arange(5, at::kFloat);
|
||||
|
||||
// PyTorch throws for empty index list
|
||||
ASSERT_THROW(t.index(std::initializer_list<at::indexing::TensorIndex>{}),
|
||||
std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, IndexWithTensorInitializerList) {
|
||||
at::Tensor t = at::arange(5, at::kFloat);
|
||||
|
||||
at::Tensor idx = at::empty({3}, at::kLong);
|
||||
int64_t* idx_data = idx.data_ptr<int64_t>();
|
||||
idx_data[0] = 0;
|
||||
idx_data[1] = 2;
|
||||
idx_data[2] = 4;
|
||||
|
||||
at::Tensor result = t.index({idx});
|
||||
|
||||
ASSERT_EQ(result.numel(), 3);
|
||||
float* result_data = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(result_data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(result_data[1], 2.0f);
|
||||
ASSERT_FLOAT_EQ(result_data[2], 4.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, MemberIndexWithArrayRefTensorIndices) {
|
||||
if (!FLAGS_use_stride_kernel) {
|
||||
return;
|
||||
}
|
||||
at::Tensor base = at::arange(24, at::kFloat).reshape({4, 6});
|
||||
at::Tensor transposed = base.t();
|
||||
std::vector<at::indexing::TensorIndex> indices = {at::indexing::Slice(1, 5),
|
||||
at::indexing::Slice(0, 3)};
|
||||
|
||||
at::Tensor sliced = transposed.index(indices);
|
||||
|
||||
ASSERT_EQ(sliced.sizes(), c10::IntArrayRef({4, 3}));
|
||||
ASSERT_EQ(sliced.strides(), c10::IntArrayRef({1, 6}));
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, MixedSliceAndTensorIndicesThrows) {
|
||||
at::Tensor t = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 0;
|
||||
idx.data_ptr<int64_t>()[1] = 2;
|
||||
|
||||
ASSERT_THROW(t.index({at::indexing::Slice(0, 2), idx}), std::exception);
|
||||
}
|
||||
|
||||
// ======================== index_put_ tests ========================
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutInplaceWithTensor) {
|
||||
at::Tensor t = at::zeros({5}, at::kFloat);
|
||||
float* original_data_ptr = t.data_ptr<float>();
|
||||
|
||||
// Create index tensor [1, 3]
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
int64_t* idx_data = idx.data_ptr<int64_t>();
|
||||
idx_data[0] = 1;
|
||||
idx_data[1] = 3;
|
||||
|
||||
// Values to put
|
||||
at::Tensor values = at::full({2}, 99.0f, at::kFloat);
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx);
|
||||
|
||||
t.index_put_(indices, values);
|
||||
|
||||
// Verify data pointer unchanged (inplace)
|
||||
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 99.0f);
|
||||
ASSERT_FLOAT_EQ(data[2], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[3], 99.0f);
|
||||
ASSERT_FLOAT_EQ(data[4], 0.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutInplaceWithScalar) {
|
||||
at::Tensor t = at::zeros({5}, at::kFloat);
|
||||
float* original_data_ptr = t.data_ptr<float>();
|
||||
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
int64_t* idx_data = idx.data_ptr<int64_t>();
|
||||
idx_data[0] = 0;
|
||||
idx_data[1] = 4;
|
||||
|
||||
t.index_put_({idx}, at::Scalar(7.0));
|
||||
|
||||
// Verify data pointer unchanged (inplace)
|
||||
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 7.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[4], 7.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutNonInplace) {
|
||||
at::Tensor t = at::zeros({5}, at::kFloat);
|
||||
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
int64_t* idx_data = idx.data_ptr<int64_t>();
|
||||
idx_data[0] = 1;
|
||||
idx_data[1] = 3;
|
||||
|
||||
at::Tensor values = at::full({2}, 42.0f, at::kFloat);
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx);
|
||||
|
||||
at::Tensor result = t.index_put(indices, values);
|
||||
|
||||
// Original should be unchanged
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[1], 0.0f);
|
||||
|
||||
// Result should have the values
|
||||
float* rdata = result.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(rdata[1], 42.0f);
|
||||
ASSERT_FLOAT_EQ(rdata[3], 42.0f);
|
||||
}
|
||||
|
||||
// ======================= Additional index edge case tests
|
||||
// =======================
|
||||
|
||||
TEST(TensorIndexTest, IndexWithEmptyList) {
|
||||
// Test index with empty indices list (should return self)
|
||||
at::Tensor t = at::arange(5, at::kFloat);
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
|
||||
at::Tensor result = t.index(indices);
|
||||
ASSERT_EQ(result.numel(), 5);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, IndexWithMultipleIndices) {
|
||||
// Test index with multiple indices (2D indexing)
|
||||
at::Tensor t = at::arange(9, at::kFloat).reshape({3, 3});
|
||||
|
||||
at::Tensor idx0 = at::empty({2}, at::kLong);
|
||||
int64_t* idx0_data = idx0.data_ptr<int64_t>();
|
||||
idx0_data[0] = 0;
|
||||
idx0_data[1] = 1;
|
||||
|
||||
at::Tensor idx1 = at::empty({2}, at::kLong);
|
||||
int64_t* idx1_data = idx1.data_ptr<int64_t>();
|
||||
idx1_data[0] = 0;
|
||||
idx1_data[1] = 2;
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx0);
|
||||
indices.push_back(idx1);
|
||||
|
||||
at::Tensor result = t.index(indices);
|
||||
ASSERT_EQ(result.numel(), 2);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, IndexWithOptionalNone) {
|
||||
// Test index with optional None in indices
|
||||
// None means "select all" along that dimension
|
||||
at::Tensor t = at::arange(9, at::kFloat).reshape({3, 3});
|
||||
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 0;
|
||||
idx.data_ptr<int64_t>()[1] = 2;
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(::std::nullopt); // None = select all rows
|
||||
indices.push_back(idx); // [0, 2] = select columns 0 and 2
|
||||
|
||||
at::Tensor result = t.index(indices);
|
||||
// Result should be shape {3, 2} = 6 elements
|
||||
// Columns 0 and 2 from all rows: [[0,2], [3,5], [6,8]]
|
||||
ASSERT_EQ(result.numel(), 6);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, FreeIndexWithAllNoneReturnsSelf) {
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(::std::nullopt);
|
||||
indices.push_back(::std::nullopt);
|
||||
|
||||
at::Tensor result = at::index(t, indices);
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[5], 5.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, FreeIndexWithSingleLeadingTensor) {
|
||||
at::Tensor t = at::arange(9, at::kFloat).reshape({3, 3});
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 2;
|
||||
idx.data_ptr<int64_t>()[1] = 0;
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx);
|
||||
|
||||
at::Tensor result = at::index(t, indices);
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 3}));
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 6.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[1], 7.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[2], 8.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[3], 0.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[4], 1.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[5], 2.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, MixedTensorNoneFullSliceIndex) {
|
||||
at::Tensor base = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 2;
|
||||
idx.data_ptr<int64_t>()[1] = 0;
|
||||
|
||||
at::Tensor result =
|
||||
base.index({idx, at::indexing::None, at::indexing::Slice()});
|
||||
|
||||
ASSERT_EQ(result.sizes(), c10::IntArrayRef({2, 1, 4}));
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[0], 8.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[1], 9.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[2], 10.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[3], 11.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[4], 0.0f);
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[7], 3.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexTest, MixedFullSliceWithMultipleTensorIndicesThrows) {
|
||||
at::Tensor base = at::arange(12, at::kFloat).reshape({3, 4});
|
||||
at::Tensor idx0 = at::empty({2}, at::kLong);
|
||||
idx0.data_ptr<int64_t>()[0] = 0;
|
||||
idx0.data_ptr<int64_t>()[1] = 1;
|
||||
at::Tensor idx1 = at::empty({2}, at::kLong);
|
||||
idx1.data_ptr<int64_t>()[0] = 0;
|
||||
idx1.data_ptr<int64_t>()[1] = 1;
|
||||
|
||||
ASSERT_THROW(base.index({idx0, at::indexing::Slice(), idx1}), std::exception);
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutAccumulate) {
|
||||
// Test index_put_ with accumulate=true
|
||||
at::Tensor t = at::zeros({5}, at::kFloat);
|
||||
float* original_data_ptr = t.data_ptr<float>();
|
||||
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 1;
|
||||
idx.data_ptr<int64_t>()[1] = 1;
|
||||
|
||||
at::Tensor values = at::full({2}, 5.0f, at::kFloat);
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx);
|
||||
|
||||
t.index_put_(indices, values, true); // accumulate=true
|
||||
|
||||
// Verify data pointer unchanged (inplace)
|
||||
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(data[1], 10.0f); // 5 + 5 (accumulated)
|
||||
ASSERT_FLOAT_EQ(data[2], 0.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutWith2D) {
|
||||
// Test index_put_ with 2D tensor
|
||||
at::Tensor t = at::zeros({3, 3}, at::kFloat);
|
||||
float* original_data_ptr = t.data_ptr<float>();
|
||||
|
||||
at::Tensor idx0 = at::arange(2, at::kLong);
|
||||
idx0.data_ptr<int64_t>()[0] = 0;
|
||||
idx0.data_ptr<int64_t>()[1] = 1;
|
||||
at::Tensor idx1 = at::arange(2, at::kLong);
|
||||
idx1.data_ptr<int64_t>()[0] = 0;
|
||||
idx1.data_ptr<int64_t>()[1] = 1;
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx0);
|
||||
indices.push_back(idx1);
|
||||
|
||||
at::Tensor values = at::full({2}, 9.0f, at::kFloat);
|
||||
|
||||
t.index_put_(indices, values);
|
||||
|
||||
// Verify data pointer unchanged (inplace)
|
||||
ASSERT_EQ(t.data_ptr<float>(), original_data_ptr);
|
||||
|
||||
float* data = t.data_ptr<float>();
|
||||
ASSERT_FLOAT_EQ(data[0], 9.0f); // [0,0]
|
||||
ASSERT_FLOAT_EQ(data[4], 9.0f); // [1,1]
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutNonInplaceAccumulate) {
|
||||
// Test index_put with accumulate=true (non-inplace)
|
||||
at::Tensor t = at::zeros({5}, at::kFloat);
|
||||
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 1;
|
||||
idx.data_ptr<int64_t>()[1] = 1;
|
||||
at::Tensor values = at::full({2}, 3.0f, at::kFloat);
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> indices;
|
||||
indices.push_back(idx);
|
||||
|
||||
at::Tensor result = t.index_put(indices, values, true);
|
||||
|
||||
// Original unchanged
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[1], 0.0f);
|
||||
// Result has accumulated
|
||||
ASSERT_FLOAT_EQ(result.data_ptr<float>()[1], 6.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutArrayRefWithTensorValue) {
|
||||
at::Tensor t = at::zeros({5}, at::kFloat);
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 1;
|
||||
idx.data_ptr<int64_t>()[1] = 4;
|
||||
at::Tensor values = at::full({2}, 13.0f, at::kFloat);
|
||||
|
||||
std::vector<at::indexing::TensorIndex> tensor_indices = {idx};
|
||||
t.index_put_(at::ArrayRef<at::indexing::TensorIndex>(tensor_indices), values);
|
||||
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[0], 0.0f);
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[1], 13.0f);
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[2], 0.0f);
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[4], 13.0f);
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutArrayRefWithNoneValue) {
|
||||
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
||||
at::Tensor values = at::full({1, 2, 3}, 6.0f, at::kFloat);
|
||||
|
||||
t.index_put_({at::indexing::None}, values);
|
||||
|
||||
ASSERT_EQ(t.sizes(), c10::IntArrayRef({2, 3}));
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 6.0f);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutArrayRefWithTensorNoneAndSlice) {
|
||||
at::Tensor t = at::zeros({3, 4}, at::kFloat);
|
||||
at::Tensor idx = at::empty({2}, at::kLong);
|
||||
idx.data_ptr<int64_t>()[0] = 2;
|
||||
idx.data_ptr<int64_t>()[1] = 0;
|
||||
at::Tensor values = at::full({2, 1, 4}, 8.0f, at::kFloat);
|
||||
|
||||
t.index_put_({idx, at::indexing::None, at::indexing::Slice()}, values);
|
||||
|
||||
ASSERT_EQ(t.sizes(), c10::IntArrayRef({3, 4}));
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 8.0f);
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[8 + i], 8.0f);
|
||||
}
|
||||
for (int i = 4; i < 8; ++i) {
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TensorIndexPutTest, IndexPutArrayRefWithNoneScalarValue) {
|
||||
at::Tensor t = at::zeros({2, 3}, at::kFloat);
|
||||
|
||||
t.index_put_({at::indexing::None}, at::Scalar(4.0));
|
||||
|
||||
ASSERT_EQ(t.sizes(), c10::IntArrayRef({2, 3}));
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
ASSERT_FLOAT_EQ(t.data_ptr<float>()[i], 4.0f);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::Tensor::item() / at::Tensor::item<T>()
|
||||
// ============================================================
|
||||
|
||||
TEST(TensorItemTest, ItemFloat_ReturnsScalar) {
|
||||
// item() on a single-element float tensor returns an at::Scalar
|
||||
at::Tensor t = at::tensor({3.14f}, at::kFloat);
|
||||
at::Scalar s = t.item();
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 3.14f, 1e-5f);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemDouble_ReturnsScalar) {
|
||||
// item() on a single-element double tensor
|
||||
at::Tensor t = at::tensor({2.718281828}, at::kDouble);
|
||||
at::Scalar s = t.item();
|
||||
|
||||
ASSERT_NEAR(s.to<double>(), 2.718281828, 1e-9);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemInt32_ReturnsScalar) {
|
||||
// item() on a single-element int32 tensor
|
||||
at::Tensor t = at::tensor({42}, at::kInt);
|
||||
at::Scalar s = t.item();
|
||||
|
||||
ASSERT_EQ(s.to<int32_t>(), 42);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemInt64_ReturnsScalar) {
|
||||
// item() on a single-element int64 tensor
|
||||
at::Tensor t = at::tensor({static_cast<int64_t>(1234567890)}, at::kLong);
|
||||
at::Scalar s = t.item();
|
||||
|
||||
ASSERT_EQ(s.to<int64_t>(), 1234567890LL);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemTemplated_Float) {
|
||||
// item<float>() returns float directly
|
||||
at::Tensor t = at::tensor({1.5f}, at::kFloat);
|
||||
float val = t.item<float>();
|
||||
|
||||
ASSERT_FLOAT_EQ(val, 1.5f);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemTemplated_Double) {
|
||||
// item<double>() returns double directly
|
||||
at::Tensor t = at::tensor({1.0 / 3.0}, at::kDouble);
|
||||
double val = t.item<double>();
|
||||
|
||||
ASSERT_NEAR(val, 1.0 / 3.0, 1e-15);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemTemplated_Int32) {
|
||||
// item<int32_t>() on int32 tensor
|
||||
at::Tensor t = at::tensor({-7}, at::kInt);
|
||||
int32_t val = t.item<int32_t>();
|
||||
|
||||
ASSERT_EQ(val, -7);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemFromSqueezed1D) {
|
||||
// item() works on a tensor that has been reshaped to single element via
|
||||
// squeeze / indexing
|
||||
at::Tensor t = at::arange(6, at::kFloat).reshape({2, 3});
|
||||
at::Tensor elem = t[1][2]; // value = 5.0
|
||||
|
||||
ASSERT_FLOAT_EQ(elem.item<float>(), 5.0f);
|
||||
}
|
||||
|
||||
TEST(TensorItemTest, ItemOnMultiElementTensorThrows) {
|
||||
// item() on a tensor with more than one element must throw.
|
||||
at::Tensor t = at::ones({2, 3}, at::kFloat);
|
||||
ASSERT_THROW(t.item(), std::exception);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::Tensor::is_variable()
|
||||
// ============================================================
|
||||
|
||||
TEST(TensorIsVariableTest, AlwaysReturnsTrue) {
|
||||
// is_variable() is always true in the eager execution mode.
|
||||
at::Tensor t = at::ones({3, 4}, at::kFloat);
|
||||
ASSERT_TRUE(t.is_variable());
|
||||
}
|
||||
|
||||
TEST(TensorIsVariableTest, AlwaysTrueForScalarTensor) {
|
||||
at::Tensor t = at::tensor({1.0f}, at::kFloat);
|
||||
ASSERT_TRUE(t.is_variable());
|
||||
}
|
||||
|
||||
TEST(TensorIsVariableTest, AlwaysTrueFor1D) {
|
||||
at::Tensor t = at::arange(10, at::kFloat);
|
||||
ASSERT_TRUE(t.is_variable());
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::Tensor::item() — sparse tensor paths
|
||||
// ============================================================
|
||||
|
||||
TEST(TensorItemSparseTest, EmptySparseCOO_ItemReturnsZero) {
|
||||
// A sparse tensor with nnz == 0: item() must return zero (Scalar(0)).
|
||||
at::Tensor indices = at::zeros({2, 0}, at::kLong);
|
||||
at::Tensor values = at::zeros({0}, at::kFloat);
|
||||
// 1x1 empty sparse tensor
|
||||
at::Tensor sparse = at::sparse_coo_tensor(indices, values, {1, 1});
|
||||
sparse = sparse.coalesce();
|
||||
|
||||
at::Scalar s = sparse.item();
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 0.0f, 1e-6f);
|
||||
}
|
||||
|
||||
TEST(TensorItemSparseTest, CoalescedSparseCOO_SingleNonZero_ReturnsValue) {
|
||||
// 1x1 sparse COO with one non-zero at (0,0) = 5.0.
|
||||
at::Tensor indices = at::tensor({0, 0}, at::kLong).reshape({2, 1});
|
||||
at::Tensor values = at::tensor({5.0f}, at::kFloat);
|
||||
at::Tensor sparse = at::sparse_coo_tensor(indices, values, {1, 1});
|
||||
sparse = sparse.coalesce();
|
||||
|
||||
ASSERT_TRUE(sparse.is_coalesced());
|
||||
at::Scalar s = sparse.item();
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 5.0f, 1e-5f);
|
||||
}
|
||||
|
||||
TEST(TensorItemSparseTest, NonCoalescedSparseCOO_DuplicateIndices_SumsValues) {
|
||||
// Two entries both at (0,0): item() must sum them (3 + 7 = 10).
|
||||
at::Tensor indices = at::tensor({0, 0, 0, 0}, at::kLong).reshape({2, 2});
|
||||
at::Tensor values = at::tensor({3.0f, 7.0f}, at::kFloat);
|
||||
at::Tensor sparse = at::sparse_coo_tensor(indices, values, {1, 1});
|
||||
|
||||
// Do NOT coalesce — exercising the non-coalesced path.
|
||||
ASSERT_FALSE(sparse.is_coalesced());
|
||||
at::Scalar s = sparse.item();
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 10.0f, 1e-5f);
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
// Copyright (c) 2026 PaddlePaddle 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 <ATen/Functions.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
#include <ATen/native/cuda/Resize.h>
|
||||
#include <ATen/ops/_local_scalar_dense.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include "ATen/ATen.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/common/bfloat16.h"
|
||||
#include "paddle/phi/common/float16.h"
|
||||
#include "torch/all.h"
|
||||
|
||||
// ============================================================
|
||||
// Tests for at::_local_scalar_dense()
|
||||
// ============================================================
|
||||
|
||||
TEST(LocalScalarDenseTest, Float32_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({2.5f}, at::kFloat);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 2.5f, 1e-6f);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Float64_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({3.141592653589793}, at::kDouble);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_NEAR(s.to<double>(), 3.141592653589793, 1e-12);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Float16_ReturnsCorrectValue) {
|
||||
// Create FP16 tensor from float, then read back via _local_scalar_dense.
|
||||
at::Tensor t = at::tensor({1.5f}, at::kHalf);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
// Float16 has ~3 significant decimal digits.
|
||||
ASSERT_NEAR(s.to<float>(), 1.5f, 1e-2f);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, BFloat16_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({1.0f}, at::kBFloat16);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 1.0f, 1e-2f);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Int8_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({static_cast<int8_t>(-7)}, at::kChar);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_EQ(s.to<int8_t>(), static_cast<int8_t>(-7));
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Int16_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({static_cast<int16_t>(300)}, at::kShort);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_EQ(s.to<int16_t>(), static_cast<int16_t>(300));
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Int32_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({42}, at::kInt);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_EQ(s.to<int32_t>(), 42);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Int64_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({static_cast<int64_t>(9876543210LL)}, at::kLong);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_EQ(s.to<int64_t>(), 9876543210LL);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, UInt8_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({static_cast<uint8_t>(255)}, at::kByte);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_EQ(s.to<uint8_t>(), static_cast<uint8_t>(255));
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Bool_True_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({true}, at::kBool);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_TRUE(s.to<bool>());
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, Bool_False_ReturnsCorrectValue) {
|
||||
at::Tensor t = at::tensor({false}, at::kBool);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_FALSE(s.to<bool>());
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, NegativeValue_Float) {
|
||||
at::Tensor t = at::tensor({-99.0f}, at::kFloat);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), -99.0f, 1e-5f);
|
||||
}
|
||||
|
||||
TEST(LocalScalarDenseTest, ZeroValue_Int32) {
|
||||
at::Tensor t = at::tensor({0}, at::kInt);
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_EQ(s.to<int32_t>(), 0);
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
TEST(LocalScalarDenseTest, GPU_Float32_ReturnsCorrectValue) {
|
||||
// _local_scalar_dense must copy to CPU when the tensor is on GPU.
|
||||
at::Tensor t = at::tensor(
|
||||
{7.0f},
|
||||
at::TensorOptions().dtype(at::kFloat).device(c10::Device(c10::kCUDA, 0)));
|
||||
at::Scalar s = at::_local_scalar_dense(t);
|
||||
|
||||
ASSERT_NEAR(s.to<float>(), 7.0f, 1e-5f);
|
||||
}
|
||||
#endif
|
||||
|
||||
TEST(LocalScalarDenseTest, EmptyTensor_ThrowsCheck) {
|
||||
// Passing an empty tensor should trigger PD_CHECK in the implementation.
|
||||
at::Tensor t = at::empty({0}, at::kFloat);
|
||||
ASSERT_THROW(at::_local_scalar_dense(t), std::exception);
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user