Skip to content

Fix computeAt reorder issue. #242

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 28, 2020
Merged

Fix computeAt reorder issue. #242

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a28b20e
Fix a couple issues. Getting reduction4 failure though.
csarofeen Jul 28, 2020
2e48137
Adjust test error tolerance
naoyam Jul 28, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 55 additions & 19 deletions test/cpp/jit/test_gpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1277,15 +1277,13 @@ int ceilDiv_(int a, int b) {

void testGPU_FusionAdvancedComputeAt() {
// Case 1
/*
* tv1 = tv0 * 0.5
* tv2 = tv1 * -1
* tv3 = tv1 + 3
* tv4 = tv1 * 2
* tv5 = tv3 + tv2
* tv6 = tv5 + tv4
* tv7 = tv1 + tv4
*/
// tv1 = tv0 * 0.5
// tv2 = tv1 * -1
// tv3 = tv1 + 3
// tv4 = tv1 * 2
// tv5 = tv3 + tv2
// tv6 = tv5 + tv4
// tv7 = tv1 + tv4
{
Fusion fusion;
FusionGuard fg(&fusion);
Expand Down Expand Up @@ -1355,14 +1353,12 @@ void testGPU_FusionAdvancedComputeAt() {
}

// Case 2
/*
* tv1 = tv0 * -1
* tv2 = tv0 + 3
* tv3 = tv0 * 2
* tv4 = tv2 + tv1
* tv5 = tv4 + tv3
* tv6 = tv5 + tv3
*/
// tv1 = tv0 * -1
// tv2 = tv0 + 3
// tv3 = tv0 * 2
// tv4 = tv2 + tv1
// tv5 = tv4 + tv3
// tv6 = tv5 + tv3
{
Fusion fusion;
FusionGuard fg(&fusion);
Expand Down Expand Up @@ -1550,6 +1546,44 @@ void testGPU_FusionAdvancedComputeAt() {

TORCH_CHECK(at::allclose(outputs[0], t6), actual_kernel.str());
}

// Case 5
// tv2 = tv0 + 2.0
// tv3 = tv1 * tv2
{
Fusion fusion;
FusionGuard fg(&fusion);

// Set up your input tensor views
TensorView* tv0 = makeDummyTensor(2);
fusion.addInput(tv0);
TensorView* tv1 = makeDummyTensor(2);
fusion.addInput(tv1);
TensorView* tv2 = add(tv0, new Float(2.0));
TensorView* tv3 = mul(tv1, tv2);
fusion.addOutput(tv3);

tv3->merge(0);
tv3->split(-1, 8);
tv3->split(-1, 4);

tv2->computeAt(tv3, 1);
tv2->split(-1, 4); // Kernel will break without this split
tv3->axis(0)->parallelize(ParallelType::BIDx);

auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
at::Tensor t0 = at::randn({63, 65}, options);
at::Tensor t1 = at::rand_like(t0, options);

auto t2 = t0.add(2.0);
auto t3 = t1.mul(t2);

torch::jit::fuser::cuda::FusionExecutor fe;
fe.compileFusion(&fusion);
auto outputs = fe.runFusion({t0, t1});

TORCH_CHECK(at::allclose(outputs[0], t3));
}
}

void testGPU_FusionScalarInputs() {
Expand Down Expand Up @@ -2577,7 +2611,6 @@ void testGPU_FusionReduction4() {
TensorView* tv2 = tv1->rFactor({-3});

tv0->computeAt(tv1, 1);

tv1->axis(0)->parallelize(ParallelType::BIDy);

for (auto* val : fusion.vals()) {
Expand All @@ -2600,7 +2633,10 @@ void testGPU_FusionReduction4() {
fe.runFusion({input}, {cg_output});

auto aten_output = input.sum({1});
TORCH_CHECK(aten_output.allclose(cg_output));
TORCH_CHECK(
aten_output.allclose(cg_output, 1e-5, 1e-7),
"Error of: ",
aten_output.sub(cg_output).abs().max());
}

void testGPU_FusionReduction5() {
Expand Down
5 changes: 2 additions & 3 deletions torch/csrc/jit/codegen/cuda/compute_at.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,9 +220,8 @@ unsigned int ComputeAt::forwardComputeAt_impl(
}

consumer_entry.setPassPosition(replay.second);
if ((consumer_entry.shouldSetComputeAt(replay.second) &&
consumer != consumer_) ||
(consumer == consumer_ && replay.second >= consumer_position_)) {
if (consumer_entry.shouldSetComputeAt(replay.second) &&
consumer != consumer_) {
consumer_entry.setComputeAtDomain(consumer->domain());
}

Expand Down
2 changes: 1 addition & 1 deletion torch/csrc/jit/codegen/cuda/kernel_resource_strings.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ __device__ __inline__ float uniform(unsigned int x) {

// Helper functions for Operations
static auto code_helper_funcs = R"(
__device__ int ceilDiv(const int a, const int b) {
__device__ constexpr int ceilDiv(const int a, const int b) {
return (a + b - 1) / b;
}
__device__ float clamp(const float x, const float minv, const float maxv) {
Expand Down