diff --git a/cinn/frontend/net_builder_test.cc b/cinn/frontend/net_builder_test.cc index 554b308d79..717ecb6a61 100755 --- a/cinn/frontend/net_builder_test.cc +++ b/cinn/frontend/net_builder_test.cc @@ -269,14 +269,12 @@ TEST(net_build, program_execute_gather) { scope->Var(std::string(output->id)); auto input1_tensor = scope->GetTensor(std::string(input1.id())); - // SetRandData(input1_tensor, target); + SetRandData(input1_tensor, target); float* input1_data = input1_tensor->mutable_data(target); auto input2_tensor = scope->GetTensor(std::string(input2.id())); - // SetRandData(input2_tensor, target); + SetRandInt(input2_tensor, target); int* input2_data = input2_tensor->mutable_data(target); - - memset(input1_data, 0, sizeof(float) * B * H_IN1); memset(input2_data, 0, sizeof(int) * B * H_IN2); runtime_program->Execute(); @@ -329,7 +327,7 @@ TEST(net_build, program_execute_gather_nd) { float* input1_data = input1_tensor->mutable_data(target); auto input2_tensor = scope->GetTensor(std::string(input2.id())); - SetRandData(input2_tensor, target); + SetRandInt(input2_tensor, target); int* input2_data = input2_tensor->mutable_data(target); runtime_program->Execute(); @@ -380,11 +378,11 @@ TEST(net_build, program_execute_scatter) { scope->Var(std::string(output->id)); auto input1_tensor = scope->GetTensor(std::string(input1.id())); - // SetRandData(input1_tensor, target); + SetRandData(input1_tensor, target); float* input1_data = input1_tensor->mutable_data(target); auto input2_tensor = scope->GetTensor(std::string(input2.id())); - // SetRandData(input2_tensor, target); + // SetRandInt(input2_tensor, target); int* input2_data = input2_tensor->mutable_data(target); runtime_program->Execute(); @@ -452,7 +450,7 @@ TEST(net_build, program_execute_scatter_nd) { SetRandData(input1_tensor, target); auto input2_tensor = scope->GetTensor(std::string(input2.id())); - SetRandData(input2_tensor, target); + SetRandInt(input2_tensor, target); runtime_program->Execute(); @@ -517,7 +515,7 @@ TEST(net_build, program_execute_cast) { scope->Var(std::string(output->id)); auto input_tensor = scope->GetTensor(std::string(input.id())); - SetRandData(input_tensor, target); + SetRandInt(input_tensor, target); int* input_data = input_tensor->mutable_data(target); runtime_program->Execute(); diff --git a/cinn/hlir/op/contrib/gather.cc b/cinn/hlir/op/contrib/gather.cc index 9a566f4ccd..bef7a7dd18 100644 --- a/cinn/hlir/op/contrib/gather.cc +++ b/cinn/hlir/op/contrib/gather.cc @@ -55,7 +55,7 @@ ir::Tensor Gather(const ir::Tensor &A, const ir::Tensor &B, const int &axis, con for (int i = 0; i < axis; ++i) { A_indices.push_back(indices[i]); } - A_indices.push_back(ir::Cast::Make(Int(32), B(indices))); + A_indices.push_back(B(indices)); for (size_t i = axis + 1; i < A->shape.size(); ++i) { A_indices.push_back(indices[i]); } diff --git a/cinn/utils/data_util.cc b/cinn/utils/data_util.cc index b9cbc339aa..96d01b2dc0 100644 --- a/cinn/utils/data_util.cc +++ b/cinn/utils/data_util.cc @@ -14,10 +14,11 @@ #include "cinn/utils/data_util.h" +#include "iostream" + namespace cinn { -template -void SetRandInt(hlir::framework::Tensor tensor, const common::Target& target, int seed) { +void SetRandInt(hlir::framework::Tensor tensor, const common::Target& target, int seed) { if (seed == -1) { std::random_device rd; seed = rd(); @@ -25,20 +26,21 @@ void SetRandInt(hlir::framework::Tensor tensor, const common::Target& targe std::default_random_engine engine(seed); std::uniform_int_distribution dist(1, 10); size_t num_ele = tensor->shape().numel(); - std::vector random_data(num_ele); + std::vector random_data(num_ele); for (size_t i = 0; i < num_ele; i++) { - random_data[i] = static_cast dist(engine); // All random data + random_data[i] = static_cast(dist(engine)); // All random data } - auto* data = tensor->mutable_data(target); + auto* data = tensor->mutable_data(target); #ifdef CINN_WITH_CUDA if (target == common::DefaultNVGPUTarget()) { - cudaMemcpy(data, random_data.data(), num_ele * sizeof(T), cudaMemcpyHostToDevice); + cudaMemcpy(data, random_data.data(), num_ele * sizeof(int), cudaMemcpyHostToDevice); return; } #endif CHECK(target == common::DefaultHostTarget()); std::copy(random_data.begin(), random_data.end(), data); + std::cout << "success" << std::endl; } template <> diff --git a/cinn/utils/data_util.h b/cinn/utils/data_util.h index 57ad14a916..d1d765c549 100644 --- a/cinn/utils/data_util.h +++ b/cinn/utils/data_util.h @@ -23,7 +23,6 @@ #endif namespace cinn { -template void SetRandInt(hlir::framework::Tensor tensor, const common::Target& target, int seed = -1); template