Google Git
Sign in
android / platform / external / pytorch / e69c2f85f6 / . / test / cpp / tensorexpr / test_registerizer.cpp
blob: 0b8a5db46f13a33fd776532ed480dcc6dc279464 [file] [log] [blame]
#include <gtest/gtest.h>
#include "test/cpp/tensorexpr/test_base.h"
#include "test/cpp/tensorexpr/test_utils.h"
#include "torch/csrc/jit/tensorexpr/ir_simplifier.h"
#include "torch/csrc/jit/tensorexpr/registerizer.h"
#include <iostream>
namespace torch {
namespace jit {
using namespace torch::jit::tensorexpr;
// Can replace a simple scalar access with a local variable.
TEST(Registerizer, RegisterizerSimple) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt =
Block::make({Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0}, Add::make(Load::make(a, {0}, 1), x), 1)}))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* A_1 = x + A_1;
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Won't do replacement of a loop access.
TEST(Registerizer, RegisterizerLoop) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
VarHandle x("x", kInt);
Stmt* stmt =
Block::make({Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {x}, Add::make(Load::make(a, {x}, 1), x), 1)}))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[x] = (A[x]) + x;
* }
*/
// No change.
stmt = registerize(stmt);
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[x] = (A[x]) + x;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK-NOT: int
# CHECK: A[0] = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A_
# CHECK: A[x] =
# CHECK-NOT: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Won't replace even if the load is a fixed scalar, since the store could
// invalidate it.
TEST(Registerizer, RegisterizerLoopFixedLoad) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt =
Block::make({Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {x}, Add::make(Load::make(a, {0}, 1), x), 1)}))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[x] = (A[0]) + x;
* }
*/
// No change.
stmt = registerize(stmt);
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[x] = (A[0]) + x;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK-NOT: int
# CHECK: A[0] = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A_
# CHECK: A[x] =
# CHECK-NOT: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// We can registerize accesses that occur entirely within inner scopes, even if
// they depend on the loop var.
TEST(Registerizer, RegisterizerLoopInternal) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Block::make(
{Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), x), 1),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), x), 1)}))});
/*
* for (int x = 0; x < 10; x++) {
* A[x] = (A[x]) + x;
* A[x] = (A[x]) + x;
* }
*/
stmt = registerize(stmt);
/*
* for (int x = 0; x < 10; x++) {
* int A_1 = A[x];
* A_1 = A_1 + x;
* A_1 = A_1 + x;
* A[x] = A_1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: for (int x = 0; x < 10; x++)
# CHECK: int A_1 = A[x];
# CHECK: A_1 = A_1 + x;
# CHECK: A_1 = A_1 + x;
# CHECK: A[x] = A_1;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// An access can be overlapped by another read in the same Expr. In this case
// B[z] and B[y] overlap and prevent registerization of both accesses.
TEST(Registerizer, RegisterizerLoopInternalLoadOverlap) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
VarHandle z("z", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Store::make(
a,
{x},
Add::make(Load::make(b, {y}, 1), Load::make(b, {z}, 1)),
1))});
/*
* for (int x = 0; x < 10; x++) {
* A[x] = (B[y]) + (B[z]);
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
TEST(Registerizer, RegisterizerLoopInternalRepeated) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {1}, 1), x), 1),
Store::make(a, {0}, Add::make(Load::make(a, {1}, 1), x), 1)})),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {1}, 1), x), 1),
Store::make(a, {0}, Add::make(Load::make(a, {1}, 1), x), 1)}))
});
/*
* for (int x = 0; x < 10; x++) {
* A[0] = x + (A[1]);
* A[0] = x + (A[1]);
* }
* for (int x = 0; x < 10; x++) {
* A[0] = x + (A[1]);
* A[0] = x + (A[1]);
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[1];
* int A_2 = A[0];
* for (int x = 0; x < 10; x++) {
* A_2 = x + A_1;
* A_2 = x + A_1;
* }
* for (int x = 0; x < 10; x++) {
* A_2 = x + A_1;
* A_2 = x + A_1;
* }
* A[0] = A_2;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[1];
# CHECK: int A_2 = A[0];
# CHECK: for (int x = 0; x < 10; x++)
# CHECK: A_2 = x + A_1;
# CHECK: A_2 = x + A_1;
# CHECK: }
# CHECK: for (int x = 0; x < 10; x++)
# CHECK: A_2 = x + A_1;
# CHECK: A_2 = x + A_1;
# CHECK: }
# CHECK-NOT: A[1]
# CHECK: A[0] = A_2;
# CHECK-NOT: A[1]
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
TEST(Registerizer, RegisterizerLoopInternalRepeatedOverlapLoopVar) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {x}, 1), x), 1),
Store::make(a, {0}, Add::make(Load::make(a, {x}, 1), x), 1)})),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {x}, 1), x), 1),
Store::make(a, {0}, Add::make(Load::make(a, {x}, 1), x), 1)}))
});
/*
* for (int x = 0; x < 10; x++) {
* A[0] = (A[x]) + x;
* A[0] = (A[x]) + x;
* }
* for (int x = 0; x < 10; x++) {
* A[0] = (A[x]) + x;
* A[0] = (A[x]) + x;
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
TEST(Registerizer, RegisterizerLoopInternalRepeatedOverlapOther) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(x, Load::make(a, {y}, 1)), 1),
Store::make(a, {0}, Add::make(x, Load::make(a, {y}, 1)), 1)})),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(x, Load::make(a, {y}, 1)), 1),
Store::make(a, {0}, Add::make(x, Load::make(a, {y}, 1)), 1)}))
});
/*
* for (int x = 0; x < 10; x++) {
* A[0] = (A[x]) + x;
* A[0] = (A[x]) + x;
* }
* for (int x = 0; x < 10; x++) {
* A[0] = (A[x]) + x;
* A[0] = (A[x]) + x;
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// Will registerize multiple accesses of different items of the same buffer.
TEST(Registerizer, RegisterizerMultiVar) {
KernelScope kernel_scope;
BufHandle a("A", {2}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({
Store::make(a, {0}, 0, 1),
Store::make(a, {1}, 0, 1),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x), 1),
Store::make(a, {1}, Sub::make(Load::make(a, {1}, 1), x), 1)})),
});
/*
* A[0] = 0;
* A[1] = 0;
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* A[1] = (A[1]) - x;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* int A_2 = 0;
* for (int x = 0; x < 10; x++) {
* A_2 = x + A_2;
* A_1 = A_1 - x;
* }
* A[1] = A_2;
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: int A_2 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A_2 =
# CHECK: A[1] = A_2
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Will registerize the valid accesses while skipping invalid replacements.
TEST(Registerizer, RegisterizerVariableLoad) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
VarHandle x2("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 0, 1),
For::make(x, 0, 10, Store::make(b, {x}, x, 1)),
For::make(
x2,
0,
10,
Block::make({Store::make(
a,
{0},
Add::make(Load::make(a, {0}, 1), Load::make(b, {x2}, 1)),
1)}))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = x;
* }
* for (int x_1 = 0; x_1 < 10; x_1++) {
* A[0] = (A[0]) + (B[x_1]);
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = x;
* }
* for (int x_1 = 0; x_1 < 10; x_1++) {
* A_1 = A_1 + (B[x_1]);
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK: B[x] = x
# CHECK: for (int x_1 = 0; x_1 < 10; x_1++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize variable accesses so long as the variable does not change.
TEST(Registerizer, RegisterizerSymbolicIndices) {
KernelScope kernel_scope;
VarHandle i("i", kInt);
VarHandle N("N", kInt);
BufHandle a("A", {N}, kInt);
VarHandle x("x", kInt);
Stmt* stmt =
Block::make({Store::make(a, {i}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {i}, Add::make(Load::make(a, {i}, 1), x), 1)}))});
/*
* A[i] = 0;
* for (int x = 0; x < 10; x++) {
* A[i] = (A[i]) + x;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* A_1 = x + A_1;
* }
* A[i] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[i] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize accesses dependent on multiple loop vars.
TEST(Registerizer, RegisterizerMultiLoop) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
For::make(
y,
0,
10,
Block::make({Store::make(
a,
{0},
Mul::make(Add::make(Load::make(a, {0}, 1), x), y),
1)})))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* for (int y = 0; y < 10; y++) {
* A[0] = x * y + (A[0]) * y;
* }
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* for (int y = 0; y < 10; y++) {
* A_1 = x * y + y * A_1;
* }
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK: for (int y = 0; y < 10; y++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize correctly if scalars already exist in the program.
TEST(Registerizer, RegisterizerRepeated) {
KernelScope kernel_scope;
BufHandle a("A", {2}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({
Store::make(a, {0}, 0, 1),
Store::make(a, {1}, 0, 1),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x), 1),
Store::make(a, {1}, Sub::make(Load::make(a, {1}, 1), x), 1)})),
});
// Registerize manually to make sure we only replace a single target.
{
registerizer::RegisterizerAnalysis analysis;
stmt->accept(&analysis);
auto candidates = analysis.getCandidates();
ASSERT_EQ(candidates.size(), 2);
candidates.pop_back();
registerizer::RegisterizerReplacer replacer(candidates);
stmt = stmt->accept_mutator(&replacer);
}
// Re-analyze and replace the second target.
{
registerizer::RegisterizerAnalysis analysis;
stmt->accept(&analysis);
auto candidates = analysis.getCandidates();
ASSERT_EQ(candidates.size(), 1);
registerizer::RegisterizerReplacer replacer(candidates);
stmt = stmt->accept_mutator(&replacer);
}
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: int A_1_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A_1_1 =
# CHECK: A[1] = A_1_1;
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize the load of A.
TEST(Registerizer, RegisterizerNoLoads) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 0, 1),
For::make(
x, 0, 10, Block::make({Store::make(a, {0}, Add::make(x, 1), 1)}))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[0] = x + 1;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* A_1 = x + 1;
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize the load of A but not the store of B.
TEST(Registerizer, RegisterizerNoRepeatedStores) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
Stmt* stmt =
Block::make({Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
b, {x}, Add::make(Load::make(a, {0}, 1), x), 1)}))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = (A[0]) + x;
* }
*/
stmt = registerize(stmt);
// TODO: its unnecessary to reorder the initializer of A[0], but it's not
// actually worse so lets not worry for now.
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = x + A_1;
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A_
# CHECK: B[x] =
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Won't registerize if there are multiple accesses which may overlap.
TEST(Registerizer, RegisterizerMultiVarOverlap) {
KernelScope kernel_scope;
BufHandle a("A", {2}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({
Store::make(a, {0}, 0, 1),
Store::make(a, {1}, 0, 1),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {x}, Add::make(Load::make(a, {0}, 1), x), 1),
Store::make(
a, {x + 1}, Sub::make(Load::make(a, {1}, 1), x), 1)})),
});
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
TEST(Registerizer, RegisterizerAllocs) {
KernelScope kernel_scope;
BufHandle a("A", {2}, kInt);
BufHandle b("B", {1}, kInt);
BufHandle c("C", {1}, kInt);
VarHandle x("x", kInt);
VarHandle b_(b.node()->base_handle());
Stmt* stmt = Block::make(
{Allocate::make(b_, kInt, {Load::make(c, {0}, 1)}),
Store::make(a, {0}, Load::make(c, {0}, 1), 1),
Store::make(b, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make(
{Store::make(b, {0}, Add::make(Load::make(b, {0}, 1), x), 1),
Store::make(a, {0}, Load::make(c, {0}, 1), 1)})),
Free::make(b_)});
/*
* Allocate(B, int, {C[0]});
* A[0] = C[0];
* B[0] = 0;
* for (int x = 0; x < 10; x++) {
* B[0] = (B[0]) + x;
* A[0] = C[0];
* }
* Free(B);
*/
stmt = registerize(stmt);
/*
* int C_1 = C[0];
* Allocate(B, int, {C_});
* int A_1 = C_1;
* int B_1 = 0;
* for (int x = 0; x < 10; x++) {
* B_1 = B_1 + x;
* A_1 = C_1;
* }
* B[0] = B_1;
* A[0] = A_1;
* Free(B);
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int C_1 = C[0];
# CHECK: Allocate(B
# CHECK: int A_1 = C_1;
# CHECK: int B_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK: B_1 =
# CHECK: A_1 = C_
# CHECK: B[0] = B_1;
# CHECK: A[0] = A_1;
# CHECK: Free(B)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
TEST(Registerizer, RegisterizerNoInitializer) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x), 1)}))});
/*
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[0];
* for (int x = 0; x < 10; x++) {
* A_1 = x + A_1;
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[0];
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
TEST(Registerizer, RegisterizerNoInitializerLoopVar) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Block::make(
{Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), x), 1)}))});
/*
* for (int x = 0; x < 10; x++) {
* A[x] = (A[x]) + x;
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
TEST(Registerizer, RegisterizerLoadThenStore) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
BufHandle b("B", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Block::make({Store::make(b, {0}, Add::make(Load::make(a, {0}, 1), x), 1),
Store::make(a, {0}, Load::make(b, {0}, 1), 1)}))});
/*
* for (int x = 0; x < 10; x++) {
* B[0] = (A[0]) + x;
* A[0] = B[0];
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[0];
* int B_1 = B[0];
* for (int x = 0; x < 10; x++) {
* B_1 = x + A_1;
* A_1 = B_1;
* }
* B[0] = B_1;
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[0];
# CHECK: int B_1 = B[0];
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: B[
# CHECK: B_1 =
# CHECK-NOT: A[
# CHECK: A_1 = B_
# CHECK: B[0] = B_
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
TEST(Registerizer, RegisterizerParallelized) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
LoopOptions loopOpts;
loopOpts.set_gpu_block_index(0);
Stmt* stmt =
Block::make({Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0}, Add::make(Load::make(a, {0}, 1), x), 1)}),
loopOpts)});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
*/
ASSERT_THROWS_WITH(
registerize(stmt),
"Registerization must occur after parallelism flattening");
}
// Should be able to registerize this since the scalar would exist before the
// branch.
TEST(Registerizer, RegisterizerConditionAfter) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {x}, Load::make(b, {x}, 1), 1),
Store::make(c, {x}, Load::make(a, {x}, 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr)});
/*
* A[x] = B[x];
* C[x] = A[x];
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = B[x];
* C[x] = A_1;
* if (x<5 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = B[x];
# CHECK: C[x] = A_1;
# CHECK: if (
# CHECK: A_1 = A_1 + 1;
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Should be able to registerize this since the scalar exists in the same form
// after the branch and there is no overlap.
TEST(Registerizer, RegisterizerConditionBefore) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr),
Store::make(a, {x}, Load::make(b, {x}, 1), 1),
Store::make(c, {x}, Load::make(a, {x}, 1), 1)});
/*
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
* A[x] = B[x];
* C[x] = A[x];
*/
stmt = registerize(stmt);
/*
* int A_ 1 = A[x];
* if (x<5 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* A_1 = B[x];
* C[x] = A_1;
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: if (
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A_1 = B[x];
# CHECK: C[x] = A_1;
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Should be able to registerize this as the combination of the two above rules.
TEST(Registerizer, RegisterizerConditionInside) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {x}, Load::make(b, {x}, 1), 1),
Store::make(c, {x}, Load::make(a, {x}, 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr),
Store::make(b, {x}, Load::make(a, {x}, 1), 1),
Store::make(a, {x}, Load::make(c, {x}, 1), 1)});
/*
* A[x] = B[x];
* C[x] = A[x];
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
* B[x] = A[x];
* A[x] = C[x];
*/
stmt = registerize(stmt);
/*
* int A_1 = B[x];
* C[x] = A_1;
* if (x<5 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* B[x] = A_1;
* A_1 = C[x];
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = B[x];
# CHECK: C[x] = A_1;
# CHECK: if (
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: B[x] = A_1;
# CHECK: A_1 = C[x];
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// An example where an access is cut by an overlapping access inside a
// condition, and both sides are large enough to be registerized but cannot be
// because there is no safe place to put the initializer or finalizer.
TEST(Registerizer, RegisterizerConditionInsideOverlap1) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {x}, Load::make(b, {x}, 1), 1),
Store::make(c, {x}, Load::make(a, {x}, 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make({
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
Store::make(a, {0}, 3, 1),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
}),
nullptr),
Store::make(b, {x}, Load::make(a, {x}, 1), 1),
Store::make(a, {x}, Load::make(c, {x}, 1), 1)});
/*
* A[x] = B[x];
* C[x] = A[x];
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* A[0] = 3;
* A[x] = (A[x]) + 1;
* }
* B[x] = A[x];
* A[x] = C[x];
*/
// The A[0] store overlaps, A[x] cutting the region that can be registerized
// into two groups.
// Each group has 2 loads and 2 stores however, so we could registerize it,
// but the first group would need to be finalized inside the condition block,
// the second would need to be initialized inside the condition block. There's
// no safe place to put these that's visible to the other uses in the group
// and so neither registerization is possible.
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// Same as the above, but the access group before the condition (and after the
// condition) are large enough to be registerized without needing the access
// from the loop. Registerization occurs but does not include any accesses in
// the condition, and the first group must be finalized before the Cond, the
// second initialized after it.
TEST(Registerizer, RegisterizerConditionInsideOverlap2) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {x}, Load::make(b, {x}, 1), 1),
Store::make(a, {x}, Load::make(b, {x + 1}, 1), 1),
Store::make(c, {x}, Load::make(a, {x}, 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make({
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
Store::make(a, {0}, 3, 1),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
}),
nullptr),
Store::make(b, {x}, Load::make(a, {x}, 1), 1),
Store::make(b, {x + 1}, Load::make(a, {x}, 1), 1),
Store::make(a, {x}, Load::make(c, {x}, 1), 1)});
/*
* A[x] = B[x];
* A[x] = B[x + 1];
* C[x] = A[x];
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* A[0] = 3;
* A[x] = (A[x]) + 1;
* }
* B[x] = A[x];
* B[x + 1] = A[x];
* A[x] = C[x];
*/
stmt = registerize(stmt);
/*
* int A_1 = B[x]; // A_1 initializer
* A_1 = B[x + 1]; //
* C[x] = A_1; //
* A[x] = A_1; // A_1 finalizer
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* A[0] = 3;
* A[x] = (A[x]) + 1;
* }
* int A_2 = A[x]; // A_2 initialier
* B[x] = A_2; //
* B[x + 1] = A_2; //
* A_2 = C[x]; //
* A[x] = A_2; // A_2 finalizer
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = B[x];
# CHECK: A_1 = B[x + 1];
# CHECK: C[x] = A_1;
# CHECK: A[x] = A_1;
# CHECK: if (
# CHECK-NOT: A_1 = A_1 + 1;
# CHECK: A[x] = (A[x]
# CHECK: A[0] =
# CHECK: A[x] = (A[x]
# CHECK: }
# CHECK: int A_2 = A[x];
# CHECK: B[x] = A_2;
# CHECK: B[x + 1] = A_2;
# CHECK: A_2 = C[x];
# CHECK: A[x] = A_2;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// When accesses are within conditional blocks they are not visible to the wider
// program, because we don't know if the branch would be taken and if it isn't
// the accesses in it don't need to be valid (think size checks on the index).
// In this case the accesses cannot be registerized.
TEST(Registerizer, RegisterizerConditionHidden) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kGT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr)});
/*
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
* if (x>5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// But... if the same access is found in a non conditional scope, that means
// that that access is valid in the higher scope (or at least if its not it's
// the user's fault). It "unhides" the conditional accesses, allowing
// registerization to occur.
TEST(Registerizer, RegisterizerConditionUnhidden) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kGT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
nullptr)});
/*
* if (x<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
* A[x] = (A[x]) + 1; <-- this is doing the unhiding.
* if (x>5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[x];
* if (x<5 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* A_1 = A_1 + 1;
* if (x>5 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: if (x<5
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A_1 = A_1 + 1;
# CHECK: if (x>5
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize a load that occurs in the condition of a Cond.
TEST(Registerizer, RegisterizerCondCondition) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {x}, Load::make(b, {x}, 1), 1),
Store::make(c, {x}, Load::make(a, {x}, 1), 1),
Cond::make(
CompareSelect::make(
Load::make(a, {x}, 1), 5, CompareSelectOperation::kLT),
Store::make(c, {x}, Add::make(Load::make(c, {x}, 1), 1), 1),
nullptr)});
/*
* A[x] = B[x];
* C[x] = A[x];
* if ((A[x])<5 ? 1 : 0) {
* C[x] = (C[x]) + 1;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = B[x];
* int C_1 = A_1;
* if (A_1<5 ? 1 : 0) {
* C_1 = C_1 + 1;
* }
* C[x] = C_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = B[x];
# CHECK: int C_1 = A_1;
# CHECK: if (A_1<5
# CHECK: C_1 = C_1 + 1;
# CHECK: C[x] = C_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Appearing in the condition of a Cond makes it visible to the enclosing scope,
// and so we can registerize internal usages.
TEST(Registerizer, RegisterizerCondConditionUnhidden) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Cond::make(
CompareSelect::make(
Load::make(a, {x}, 1), 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 1), 1),
Store::make(a, {x}, Add::make(Load::make(a, {x}, 1), 10), 1))});
/*
* if ((A[x])<5 ? 1 : 0) {
* A[x] = (A[x]) + 1;
* } else {
* A[x] = (A[x]) + 10;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[x];
* if (A_1<5 ? 1 : 0) {
* A_1 = A_1 + 1;
* } else {
* A_1 = A_1 + 10;
* }
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: if (A_1<5
# CHECK: A_1 = A_1 + 1;
# CHECK: } else {
# CHECK: A_1 = A_1 + 10;
# CHECK: }
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Conditional hiding also works for IfThenElse exprs.
TEST(Registerizer, RegisterizerIfThenElseHidden) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{Store::make(
b,
{y},
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), 1),
Add::make(Load::make(a, {x + 1}, 1), 2)),
1),
Store::make(
b,
{y + 1},
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), 1),
Add::make(Load::make(a, {x + 1}, 1), 2)),
1)});
/*
* B[y] = IfThenElse(x<5 ? 1 : 0, (A[x]) + 1, (A[x + 1]) + 2);
* B[y + 1] = IfThenElse(x<5 ? 1 : 0, (A[x]) + 1, (A[x + 1]) + 2);
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// Conditional unhiding also works for IfThenElse exprs.
TEST(Registerizer, RegisterizerIfThenElseUnhidden) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({
Store::make(a, {x}, 0, 1),
Store::make(
b,
{y},
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), 1),
Add::make(Load::make(a, {x + 1}, 1), 2)),
1),
Store::make(
b,
{y + 1},
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), 1),
Add::make(Load::make(a, {x + 1}, 1), 2)),
1),
});
/*
* A[x] = 0;
* B[y] = IfThenElse(x<5 ? 1 : 0, (A[x]) + 1, (A[x + 1]) + 2);
* B[y + 1] = IfThenElse(x<5 ? 1 : 0, (A[x]) + 1, (A[x + 1]) + 2);
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* B[y] = IfThenElse(x<5 ? 1 : 0, A_1 + 1, (A[x + 1]) + 2);
* B[y + 1] = IfThenElse(x<5 ? 1 : 0, A_1 + 1, (A[x + 1]) + 2);
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: B[y] = IfThenElse(x<5 ? 1 : 0, A_1 + 1, (A[x + 1]) + 2);
# CHECK: B[y + 1] = IfThenElse(x<5 ? 1 : 0, A_1 + 1, (A[x + 1]) + 2);
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Nested IfThenElse exprs can't promote to higher level scopes.
TEST(Registerizer, RegisterizerIfThenElseNested) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
BufHandle d("D", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Store::make(
a,
{x},
IfThenElse::make(
CompareSelect::make(x, 3, CompareSelectOperation::kLT),
IfThenElse::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Load::make(d, {x}, 1),
Load::make(b, {x}, 1)),
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kEQ),
Load::make(c, {x}, 1),
Load::make(d, {x}, 1))),
1)});
/*
* A[x] = IfThenElse(x<3 ? 1 : 0,
* IfThenElse(x==2 ? 1 : 0, D[x], B[x]),
* IfThenElse(x==5 ? 1 : 0, C[x], D[x]));
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// Cannot registerize an access completely contained within an IfThenElse
// branch, since it is not a Stmt and cannot hold variable definitions. We need
// to check that we don't promote the initializer/finalizer to the enclosing
// Block.
TEST(Registerizer, RegisterizerIfThenElseInternal) {
KernelScope kernel_scope;
// Making these floats so they don't get simplified to a single access.
BufHandle a("A", {5}, kFloat);
BufHandle b("B", {5}, kFloat);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Store::make(
a,
{x},
IfThenElse::make(
CompareSelect::make(x, 3, CompareSelectOperation::kLT),
Add::make(Load::make(b, {x}, 1), Load::make(b, {x}, 1)),
Load::make(b, {x}, 1)),
1)});
/*
* A[x] = IfThenElse(x<3 ? 1 : 0, (B[x]) + (B[x]), B[x]);
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
// If this was a Cond instead of an IfThenElse then we could registerize the
// two accesses to B[x] in the True branch.
// Actually lets verify that.
stmt = Block::make({Cond::make(
CompareSelect::make(x, 3, CompareSelectOperation::kLT),
Store::make(
a, {x}, Add::make(Load::make(b, {x}, 1), Load::make(b, {x}, 1)), 1),
Store::make(a, {x}, Load::make(b, {x}, 1), 1))});
/*
* if (x<3 ? 1 : 0) {
* A[x] = (B[x]) + (B[x]);
* } else {
* A[x] = B[x];
* }
*/
stmt = registerize(stmt);
/*
* if (x<3 ? 1 : 0) {
* float B_1 = B[x];
* A[x] = B_1 + B_1;
* } else {
* A[x] = B[x];
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK-NOT: int
# CHECK-NOT: float
# CHECK: if (x<3
# CHECK: float B_1 =
# CHECK: A[x] = B_1 + B_1
# CHECK: } else {
# CHECK: A[x] = B[x]
# CHECK: }
# CHECK-NOT: A[x]
# CHECK-NOT: B[x])IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize a load that occurs in the condition of an IfThenElse;
TEST(Registerizer, RegisterizerIfThenElseCondition) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {x}, Load::make(a, {x}, 1), 1),
Store::make(
a,
{x},
IfThenElse::make(
CompareSelect::make(
Load::make(a, {x}, 1), 5, CompareSelectOperation::kLT),
Load::make(b, {0}, 1),
Load::make(c, {0}, 1)),
1)});
/*
* A[x] = A[x]; <---- just here so there are enough accesses to combine.
* A[x] = IfThenElse((A[x])<5 ? 1 : 0, B[0], C[0]);
*/
stmt = registerize(stmt);
/*
* int A_1 = A[x];
* A_1 = A_1;
* A_1 = IfThenElse(A_1<5 ? 1 : 0, B[0], C[0]);
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: A_1 = IfThenElse(A_1<5 ? 1 : 0, B[0], C[0]);
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Appearing in the condition of a Cond makes it visible to the enclosing scope,
// and so we can registerize internal usages.
TEST(Registerizer, RegisterizerIfThenElseConditionUnhidden) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Store::make(
b,
{x},
IfThenElse::make(
CompareSelect::make(
Load::make(a, {x}, 1), 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), 1),
Add::make(Load::make(a, {x}, 1), 10)),
1)});
/*
* B[x] = IfThenElse((A[x])<5 ? 1 : 0, (A[x]) + 1, (A[x]) + 10);
*/
stmt = registerize(stmt);
/*
* int A_1 = A[x];
* B[x] = IfThenElse(A_1<5 ? 1 : 0, A_1 + 1, A_1 + 10);
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: B[x] = IfThenElse(A_1<5 ? 1 : 0, A_1 + 1, A_1 + 10);)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Cannot promote accesses internal to IfThenElse branches even if the enclosing
// scope if conditional.
TEST(Registerizer, RegisterizerConditionBranchOnly) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Block::make({
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(
a,
{x},
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), x),
Add::make(Load::make(a, {x - 5}, 1), x)),
1),
Store::make(
a,
{x - 5},
IfThenElse::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Add::make(Load::make(a, {x}, 1), x),
Add::make(Load::make(a, {x - 5}, 1), x)),
1)),
}))});
std::ostringstream before;
before << *stmt;
/* for (int x = 0; x < 10; x++) {
* if (x<5 ? 1 : 0) {
* A[x] = IfThenElse(x<5 ? 1 : 0, (A[x]) + x, (A[x - 5]) + x);
* } else {
* A[x - 5] = IfThenElse(x<5 ? 1 : 0, (A[x]) + x, (A[x - 5]) + x);
* }
* }
*/
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// We can registerize an IfThenElse that appears in the condition branch of a
// Cond. This is a weird but valid thing to do.
TEST(Registerizer, RegisterizerCondIfThenElse) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
BufHandle c("C", {5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Cond::make(
CompareSelect::make(
IfThenElse::make(
CompareSelect::make(
Load::make(a, {x}, 1), 5, CompareSelectOperation::kLT),
Load::make(a, {x}, 1),
Load::make(b, {x}, 1)),
x,
CompareSelectOperation::kEQ),
Store::make(c, {x}, Add::make(Load::make(c, {x}, 1), 1), 1),
nullptr)});
/*
* if ((IfThenElse((A[x])<5 ? 1 : 0, A[x], B[x]))==x ? 1 : 0) {
* C[x] = (C[x]) + 1;
* }
*/
stmt = registerize(stmt);
// access to A can be registerized, but not B or C
/*
* int A_1 = A[x];
* if ((IfThenElse(A_1<5 ? 1 : 0, A_1, B[x]))==x ? 1 : 0) {
* C[x] = (C[x]) + 1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: if ((IfThenElse(A_1<5 ? 1 : 0, A_1, B[x]
# CHECK: C[x] = (C[x]) + 1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can registerize a conditional access in the RHS of a store unhidden by it's
// LHS, and hoist it out of a loop.
TEST(Registerizer, RegisterizerIfThenElseLoop) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = For::make(
y,
0,
10,
Store::make(
a,
{x},
IfThenElse::make(
CompareSelect::make(x, 3, CompareSelectOperation::kLT),
Load::make(a, {x}, 1),
Load::make(b, {y}, 1)),
1));
/*
* for (int y = 0; y < 10; y++) {
* A[x] = IfThenElse(x<3 ? 1 : 0, A[x], B[y]);
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[x];
* for (int y = 0; y < 10; y++) {
* A_1 = IfThenElse(x<3 ? 1 : 0, A_1, B[y]);
* }
* A[x] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[x];
# CHECK: for (
# CHECK: A_1 = IfThenElse(x<3 ? 1 : 0, A_1, B[y]);
# CHECK: }
# CHECK: A[x] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Cannot registerize if the RHS overlaps the access creating visibility.
TEST(Registerizer, RegisterizerIfThenElseLoopCut) {
KernelScope kernel_scope;
BufHandle a("A", {5}, kInt);
BufHandle b("B", {5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({For::make(
y,
0,
10,
Store::make(
a,
{x},
IfThenElse::make(
CompareSelect::make(x, 3, CompareSelectOperation::kLT),
Load::make(a, {x}, 1),
Load::make(a, {y}, 1)),
1))});
/*
* for (int y = 0; y < 10; y++) {
* A[x] = IfThenElse(x<3 ? 1 : 0, A[x], A[y]);
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// Simple case where an access is cut by an overlapping access later in the
// program, we can registerize up until the overlap.
TEST(Registerizer, RegisterizerPartialAfter) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x), 1)})),
For::make(x, 1, 10, Store::make(a, {x}, Load::make(a, {x - 1}, 1), 1))});
/*
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
* for (int x = 1; x < 10; x++) {
* A[x] = A[x - 1];
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* A_1 = A_1 + x;
* }
* A[0] = A_1;
* for (int x = 1; x < 10; x++) {
* A[x] = A[x - 1];
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (
# CHECK: A_1 = A_1 + x;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: for (
# CHECK: A[x] = A[x - 1];
# CHECK: }
# CHECK-NOT: A)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// We can registerize an access which overlaps a previous access, the
// initializer must be inserted after the previous access.
TEST(Registerizer, RegisterizerPartialBefore) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{For::make(x, 1, 10, Store::make(a, {x}, Load::make(a, {x - 1}, 1), 1)),
Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0}, Add::make(Load::make(a, {0}, 1), x), 1)}))});
/*
* for (int x = 1; x < 10; x++) {
* A[x] = A[x - 1];
* }
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
*/
stmt = registerize(stmt);
/*
* for (int x = 1; x < 10; x++) {
* A[x] = A[x - 1];
* }
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* A_1 = A_1 + x;
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK-NOT: int
# CHECK: for (
# CHECK: A[x] = A[x - 1];
# CHECK: }
# CHECK: int A_1 = 0;
# CHECK: for (
# CHECK: A_1 = A_1 + x;
# CHECK: }
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// The combination of the previous two tests, an access is cut by an overlapping
// access in both directions.
TEST(Registerizer, RegisterizerPartialInside) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x1("x1", kInt);
VarHandle x2("x2", kInt);
VarHandle x3("x3", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 2, 1),
For::make(
x1,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x1), 1)),
For::make(
x2, 1, 10, Store::make(a, {x2}, Load::make(a, {x2 - 1}, 1), 1)),
For::make(
x3,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x3), 1))});
/*
* A[0] = 2;
* for (int x1 = 0; x1 < 10; x1++) {
* A[0] = (A[0]) + x1;
* }
* for (int x2 = 1; x2 < 10; x2++) {
* A[x2] = A[x2 - 1];
* }
* for (int x3 = 0; x3 < 10; x3++) {
* A[0] = (A[0]) + x3;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 2;
* for (int x1 = 0; x1 < 10; x1++) {
* A_1 = A_1 + x1;
* }
* A[0] = A_1;
* for (int x2 = 1; x2 < 10; x2++) {
* A[x2] = A[x2 - 1];
* }
* int A_2 = A[0];
* for (int x3 = 0; x3 < 10; x3++) {
* A_2 = A_2 + x3;
* }
* A[0] = A_2;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 2;
# CHECK: for (
# CHECK: A_1 = A_1 + x1;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: for (
# CHECK: A[x2] =
# CHECK: }
# CHECK: int A_2 = A[0];
# CHECK: for (
# CHECK: A_2 = A_2 + x3;
# CHECK: }
# CHECK: A[0] = A_2;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// An element could be registerized program wide but is cut by a conditional
// access, we should break this into two scalars and write back to the buffer
// before the condition.
TEST(Registerizer, RegisterizerPartialCondition) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 2, 1),
For::make(
x,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x), 1)),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Store::make(a, {x}, Load::make(a, {x - 1}, 1), 1),
nullptr),
For::make(
x,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), x), 1))});
/*
* A[0] = 2;
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
* if (x<5 ? 1 : 0) {
* A[x] = A[x - 1];
* }
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + x;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 2;
* for (int x = 0; x < 10; x++) {
* A_1 = A_1 + x;
* }
* A[0] = A_1;
* if (x<5 ? 1 : 0) {
* A[x] = A[x - 1];
* }
* int A_2 = A[0];
* for (int x = 0; x < 10; x++) {
* A_2 = A_2 + x;
* }
* A[0] = A_2;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 2;
# CHECK: for (
# CHECK: A_1 = A_1 + x;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: if (
# CHECK: A[x] =
# CHECK: }
# CHECK: int A_2 = A[0];
# CHECK: for (
# CHECK: A_2 = A_2 + x;
# CHECK: }
# CHECK: A[0] = A_2;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Tests case where an access is cut by an internal conditional access which
// itself is registerized.
TEST(Registerizer, RegisterizerPartialConditionInternalCut) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 1, 1),
Store::make(a, {0}, 3, 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make({Store::make(a, {x}, 1, 1), Store::make(a, {x}, 3, 1)}),
nullptr),
Store::make(a, {0}, 4, 1),
Store::make(a, {0}, 6, 1)});
/*
* A[0] = 1;
* A[0] = 3;
* if (x<5 ? 1 : 0) {
* A[x] = 1;
* A[x] = 3;
* }
* A[0] = 4;
* A[0] = 6;
*/
stmt = registerize(stmt);
/*
* int A_1 = 1;
* A_1 = 3;
* A[0] = A_1;
* if (x<5 ? 1 : 0) {
* int A_2 = 1;
* A_2 = 3;
* A[x] = A_2;
* }
* int A_3 = 4;
* A_3 = 6;
* A[0] = A_3;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 1;
# CHECK: A_1 = 3
# CHECK: A[0] = A_1;
# CHECK: if (
# CHECK: int A_2 = 1;
# CHECK: A_2 = 3;
# CHECK: A[x] = A_2;
# CHECK: }
# CHECK: int A_3 = 4;
# CHECK: A_3 = 6;
# CHECK: A[0] = A_3;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// First statment in condition closes outer access, but can be registerized with
// later statements.
TEST(Registerizer, RegisterizerPartialConditionInternalStart) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, 1, 1),
Store::make(a, {0}, 3, 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make({Store::make(a, {x}, 1, 1), Store::make(a, {x}, 3, 1)}),
nullptr),
Store::make(a, {x}, 4, 1),
Store::make(a, {x}, 6, 1)});
/*
* A[0] = 1;
* A[0] = 3;
* if (x<5 ? 1 : 0) {
* A[x] = 1;
* A[x] = 3;
* }
* A[x] = 4;
* A[x] = 6;
*/
stmt = registerize(stmt);
/*
* int A_1 = 1;
* A_1 = 3;
* A[0] = A_1;
* int A_2 = A[x]; <--- must read from the input here.
* if (x<5 ? 1 : 0) {
* A_2 = 1;
* A_2 = 3;
* }
* A_2 = 4;
* A_2 = 6;
* A[x] = A_2;
*/
// TODO: I suppose we could refactor with a conditional initializier?
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 1;
# CHECK: A_1 = 3
# CHECK: A[0] = A_1;
# CHECK: int A_2 = A[x];
# CHECK: if (
# CHECK: A_2 = 1;
# CHECK: A_2 = 3;
# CHECK: }
# CHECK: A_2 = 4;
# CHECK: A_2 = 6;
# CHECK: A[x] = A_2;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// An access cuts two open overlaps and creates four scalar variables.
TEST(Registerizer, RegisterizerPartialOverlapsTwo) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Store::make(a, {1}, Load::make(a, {0}, 1), 1),
Store::make(a, {0}, Load::make(a, {1}, 1), 1),
Store::make(a, {0}, Load::make(a, {1}, 1), 1),
For::make(x, 1, 10, Store::make(a, {x}, x, 1)),
Store::make(a, {1}, Load::make(a, {0}, 1), 1),
Store::make(a, {0}, Load::make(a, {1}, 1), 1),
Store::make(a, {0}, Load::make(a, {1}, 1), 1)});
/*
* A[1] = A[0];
* A[0] = A[1];
* A[0] = A[1];
* for (int x = 1; x < 10; x++) {
* A[x] = x;
* }
* A[1] = A[0];
* A[0] = A[1];
* A[0] = A[1];
*/
stmt = registerize(stmt);
/*
* int A_1 = A[0];
* int A_2 = A_1;
* A_1 = A_2;
* A_1 = A_2;
* A[1] = A_2;
* A[0] = A_1;
* for (int x = 1; x < 10; x++) {
* A[x] = x;
* }
* int A_3 = A[0];
* int A_4 = A_3;
* A_3 = A_4;
* A_3 = A_4;
* A[1] = A_4;
* A[0] = A_3;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[0];
# CHECK: int A_2 = A_1;
# CHECK: A_1 = A_2;
# CHECK: A_1 = A_2;
# CHECK: A[1] = A_2;
# CHECK: A[0] = A_1;
# CHECK: for (
# CHECK: A[x] = x;
# CHECK: }
# CHECK: int A_3 = A[0];
# CHECK: int A_4 = A_3;
# CHECK: A_3 = A_4;
# CHECK: A_3 = A_4;
# CHECK: A[1] = A_4;
# CHECK: A[0] = A_3;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Nested blocks will automatically be flattened and do not provent
// registerization of enclosed accesses.
TEST(Registerizer, RegisterizerNestedBlocks) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 2), 1)}),
Block::make({Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 3), 1),
Block::make({Store::make(
a, {0}, Add::make(Load::make(a, {0}, 1), 4), 1)})})});
/*
* A[0] = (A[0]) + 1;
* {
* A[0] = (A[0]) + 2;
* }
* {
* A[0] = (A[0]) + 3;
* {
* A[0] = (A[0]) + 4;
* }
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[0];
* A_1 = A_1 + 1;
* A_1 = A_1 + 2;
* A_1 = A_1 + 3;
* A_1 = A_1 + 4;
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[0];
# CHECK: A_1 = A_1 + 1;
# CHECK: A_1 = A_1 + 2;
# CHECK: A_1 = A_1 + 3;
# CHECK: A_1 = A_1 + 4;
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// The access can be registerized internally to a condition, but must ensure
// that both initializer and finalizer are within the same condition.
TEST(Registerizer, RegisterizerNestedConditions) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make({Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)}),
nullptr)});
/*
* if (x<5 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* if (x==2 ? 1 : 0) {
*
* A[0] = (A[0]) + 1;
* }
* }
*/
stmt = registerize(stmt);
/*
* if (x<5 ? 1 : 0) {
* int A_1 = A[0];
* A_1 = A_1 + 1;
* if (x==2 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* A[0] = A_1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: if (x<5
# CHECK: int A_1 = A[0];
# CHECK: A_1 = A_1 + 1;
# CHECK: if (x==2
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// If an access exists outside the scope of the condition then we can lift
// nested conditional usages into the same scalar.
TEST(Registerizer, RegisterizerNestedConditionsUnhidden) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make(
{Store::make(a, {1}, 1, 1),
Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)}),
nullptr)});
/*
* A[0] = (A[0]) + 1;
* if (x<5 ? 1 : 0) {
* A[1] = 1;
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = A[0];
* A_1 = A_1 + 1;
* if (x<5 ? 1 : 0) {
* A[1] = 1;
* if (x==2 ? 1 : 0) {
* A_1 = A_1 + 1;
* }
* }
* A[0] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = A[0];
# CHECK: A_1 = A_1 + 1;
# CHECK: if (x<5
# CHECK: A[1] = 1;
# CHECK: if (x==2
# CHECK: A_1 = A_1 + 1;
# CHECK: A[0] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
TEST(Registerizer, RegisterizerNestedConditionsHiddenFirst) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make({Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)}),
nullptr)});
/*
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
* if (x<5 ? 1 : 0) {
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
stmt = registerize(stmt);
}
TEST(Registerizer, RegisterizerNestedConditionsHiddenSecond) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make({Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)}),
nullptr),
Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)});
/*
* if (x<5 ? 1 : 0) {
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
* }
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
stmt = registerize(stmt);
}
// If an access is cut by another access internal to a condition block, it still
// cuts the access.
TEST(Registerizer, RegisterizerNestedConditionsCut) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
Block::make(
{Store::make(a, {x}, 1, 1),
Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)}),
nullptr)});
/*
* A[0] = (A[0]) + 1;
* if (x<5 ? 1 : 0) {
* A[x] = 1;
* if (x==2 ? 1 : 0) {
*
* A[0] = (A[0]) + 1;
* }
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
TEST(Registerizer, RegisterizerNestedConditionLoopHidden) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr),
For::make(
x,
0,
10,
Block::make(
{Store::make(b, {x}, 0, 1),
Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
nullptr)}))});
/*
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
* for (int x = 0; x < 10; x++) {
* B[x] = 0; <-- this is only here to prevent Loop/Cond reordering.
* if (x==2 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* }
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// Three loops and four element regions, three of which should be registerized
// at different levels of the IR.
TEST(Registerizer, RegisterizerNestedConditionThreeDeep) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {4}, 0, 1),
Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kGT),
Cond::make(
CompareSelect::make(x, 3, CompareSelectOperation::kGT),
Block::make({
Cond::make(
CompareSelect::make(x, 4, CompareSelectOperation::kGT),
Block::make({
Store::make(
a, {1}, Add::make(Load::make(a, {1}, 1), 1), 1),
Store::make(
a, {2}, Add::make(Load::make(a, {2}, 1), 1), 1),
Store::make(
a, {3}, Add::make(Load::make(a, {3}, 1), 1), 1),
Store::make(
a, {4}, Add::make(Load::make(a, {4}, 1), 1), 1),
Store::make(
a, {1}, Add::make(Load::make(a, {1}, 1), 1), 1),
}),
nullptr),
Store::make(a, {2}, Add::make(Load::make(a, {2}, 1), 1), 1),
}),
nullptr),
nullptr)});
/*
* A[4] = 0;
* if (x>2 ? 1 : 0) {
* if (x>3 ? 1 : 0) {
* if (x>4 ? 1 : 0) {
* A[1] = (A[1]) + 1;
* A[2] = (A[2]) + 1;
* A[3] = (A[3]) + 1;
* A[4] = (A[4]) + 1;
* A[1] = (A[1]) + 1;
* }
* A[2] = (A[2]) + 1;
* }
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* if (x>2 ? 1 : 0) {
* if (x>3 ? 1 : 0) {
* int A_3 = A[2];
* if (x>4 ? 1 : 0) {
* int A_2 = A[1];
* A_2 = A_2 + 1;
* A_3 = A_3 + 1;
* A[3] = (A[3]) + 1;
* A_1 = A_1 + 1;
* A_2 = A_2 + 1;
* A[1] = A_2;
* }
* A_3 = A_3 + 1;
* A[2] = A_3;
* }
* }
* A[4] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: if (x>2 ? 1 : 0) {
# CHECK: if (x>3 ? 1 : 0) {
# CHECK: int A_3 = A[2];
# CHECK: if (x>4 ? 1 : 0) {
# CHECK: int A_2 = A[1];
# CHECK: A_2 = A_2 + 1;
# CHECK: A_3 = A_3 + 1;
# CHECK: A[3] = (A[3]) + 1;
# CHECK: A_1 = A_1 + 1;
# CHECK: A_2 = A_2 + 1;
# CHECK: A[1] = A_2;
# CHECK: }
# CHECK: A_3 = A_3 + 1;
# CHECK: A[2] = A_3;
# CHECK: }
# CHECK: }
# CHECK: A[4] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Can replace a simple scalar access with a local variable even when that
// variable is an outer loop var.
TEST(Registerizer, RegisterizerNestedLoopSimple) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({For::make(
y,
0,
10,
For::make(
x,
0,
10,
Block::make({Store::make(
a, {y}, Add::make(Load::make(a, {y}, 1), x), 1)})))});
/*
* for (int y = 0; y < 10; y++) {
* for (int x = 0; x < 10; x++) {
* A[y] = (A[y]) + x;
* }
* }
*/
stmt = registerize(stmt);
/*
* for (int y = 0; y < 10; y++) {
* int A_1 = A[y];
* for (int x = 0; x < 10; x++) {
* A_1 = x + A_1;
* }
* A[y] = A_1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: for (int y
# CHECK: int A_1 = A[y];
# CHECK: for (int x
# CHECK: A_1 = x + A_1;
# CHECK: }
# CHECK: A[y] = A_1;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Test the positive case of the hiddenAccess split, where an internal
// conditional access can be hoisted up through a loop to match an existing
// access in a higher scope and the two can be registerized.
TEST(Registerizer, RegisterizerHiddenAccessYes) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Block::make(
{Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
Block::make(
{Store::make(b, {x}, 0, 1),
Cond::make(
CompareSelect::make(x, 3, CompareSelectOperation::kEQ),
For::make(
y,
0,
10,
Store::make(
a,
{0},
Add::make(Load::make(a, {0}, 1), 1),
1)),
nullptr)}))}),
nullptr)});
/*
* if (x==2 ? 1 : 0) {
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = 0;
* if (x==3 ? 1 : 0) {
* for (int y = 0; y < 10; y++) {
* A[0] = (A[0]) + 1;
* }
* }
* }
* }
*/
stmt = registerize(stmt);
/*
* if (x==2 ? 1 : 0) {
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = 0;
* if (x==3 ? 1 : 0) {
* for (int y = 0; y < 10; y++) {
* A_1 = A_1 + 1;
* }
* }
* }
* A[0] = A_1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: if (x==2
# CHECK: int A_1 = 0;
# CHECK: for (int x
# CHECK: B[x] = 0;
# CHECK: if (x==3
# CHECK: for (int y
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: }
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Test the negative case of the hiddenAccess split, where the hoisted access is
// never unhidden at a higher scope and registerization occurs at the lower
// scope.
TEST(Registerizer, RegisterizerHiddenAccessNo) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Block::make({For::make(
x,
0,
10,
Block::make(
{Store::make(b, {x}, 0, 1),
Cond::make(
CompareSelect::make(x, 3, CompareSelectOperation::kEQ),
For::make(
y,
0,
10,
Store::make(
a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1)),
nullptr)}))}),
nullptr)});
/*
* if (x==2 ? 1 : 0) {
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* B[x] = 0;
* if (x==3 ? 1 : 0) {
* for (int y = 0; y < 10; y++) {
* A[0] = (A[0]) + 1;
* }
* }
* }
* }
*/
stmt = registerize(stmt);
/*
* if (x==2 ? 1 : 0) {
* for (int x = 0; x < 10; x++) {
* B[x] = 0;
* if (x==3 ? 1 : 0) {
* int A_1 = A[0];
* for (int y = 0; y < 10; y++) {
* A_1 = A_1 + 1;
* }
* A[0] = A_1;
* }
* }
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: if (x==2
# CHECK: for (int x
# CHECK: B[x] = 0;
# CHECK: if (x==3
# CHECK: int A_1 = A[0];
# CHECK: for (int y
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: }
# CHECK: }
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// In this case the conditional access must be hoisted by two loops, there are
// two accesses here one is unhidden and the other isnt. A[0] can be
// registerized but B[0] cannot.
TEST(Registerizer, RegisterizerHiddenAccessMultiLoop) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({Cond::make(
CompareSelect::make(x, 2, CompareSelectOperation::kEQ),
Block::make(
{Store::make(a, {0}, 0, 1),
For::make(
x,
0,
10,
For::make(
y,
0,
10,
Block::make({Cond::make(
CompareSelect::make(y, 3, CompareSelectOperation::kEQ),
Block::make(
{Store::make(
a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1),
Store::make(
b,
{0},
Add::make(Load::make(b, {0}, 1), 1),
1)}),
nullptr)})))}),
nullptr)});
/*
* if (x==2 ? 1 : 0) {
* A[0] = 0;
* for (int x = 0; x < 10; x++) {
* for (int y = 0; y < 10; y++) {
* if (y==3 ? 1 : 0) {
* A[0] = (A[0]) + 1;
* B[0] = (B[0]) + 1;
* }
* }
* }
* }
*/
stmt = registerize(stmt);
/*
* if (x==2 ? 1 : 0) {
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* for (int y = 0; y < 10; y++) {
* if (y==3 ? 1 : 0) {
* A_1 = A_1 + 1;
* B[0] = (B[0]) + 1;
* }
* }
* }
* A[0] = A_1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: if (x==2
# CHECK: int A_1 = 0;
# CHECK: for (int x
# CHECK: for (int y
# CHECK: if (y==3
# CHECK: A_1 = A_1 + 1;
# CHECK: B[0] = (B[0]) + 1;
# CHECK: }
# CHECK: }
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Accesses are registerized inside two conditions, but the immeidate parent is
// not a condition.
TEST(Registerizer, RegisterizerTwoConditionalLoops) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
For::make(
x,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1)),
nullptr),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kGT),
For::make(
x,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1)),
nullptr)});
/*
* if (x<5 ? 1 : 0) {
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + 1;
* }
* }
* if (x>5 ? 1 : 0) {
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + 1;
* }
* }
*/
stmt = registerize(stmt);
/*
* if (x<5 ? 1 : 0) {
* int A_1 = A[0];
* for (int x = 0; x < 10; x++) {
* A_1 = A_1 + 1;
* }
* A[0] = A_1;
* }
* if (x>5 ? 1 : 0) {
* int A_2 = A[0];
* for (int x = 0; x < 10; x++) {
* A_2 = A_2 + 1;
* }
* A[0] = A_2;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: if (x<5
# CHECK: int A_1 = A[0];
# CHECK: for (int x
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: }
# CHECK: if (x>5
# CHECK: int A_2 = A[0];
# CHECK: for (int x
# CHECK: A_2 = A_2 + 1;
# CHECK: }
# CHECK: A[0] = A_2;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Accesses are registerized inside two conditions, cut in the middle.
TEST(Registerizer, RegisterizerTwoConditionalLoopsCut) {
KernelScope kernel_scope;
BufHandle a("A", {1}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kLT),
For::make(
x,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1)),
nullptr),
For::make(x, 0, 10, Store::make(a, {x}, 1, 1)),
Cond::make(
CompareSelect::make(x, 5, CompareSelectOperation::kGT),
For::make(
x,
0,
10,
Store::make(a, {0}, Add::make(Load::make(a, {0}, 1), 1), 1)),
nullptr)});
/*
* if (x<5 ? 1 : 0) {
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + 1;
* }
* }
* for (int x = 0; x < 10; x++) {
* A[x] = 1;
* }
* if (x>5 ? 1 : 0) {
* for (int x = 0; x < 10; x++) {
* A[0] = (A[0]) + 1;
* }
* }
*/
stmt = registerize(stmt);
/*
* if (x<5 ? 1 : 0) {
* int A_1 = A[0];
* for (int x = 0; x < 10; x++) {
* A_1 = A_1 + 1;
* }
* A[0] = A_1;
* }
* for (int x = 0; x < 10; x++) {
* A[x] = 1;
* }
* if (x>5 ? 1 : 0) {
* int A_2 = A[0];
* for (int x = 0; x < 10; x++) {
* A_2 = A_2 + 1;
* }
* A[0] = A_2;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: if (x<5
# CHECK: int A_1 = A[0];
# CHECK: for (int x
# CHECK: A_1 = A_1 + 1;
# CHECK: }
# CHECK: A[0] = A_1;
# CHECK: }
# CHECK: for (int x
# CHECK: A[x] = 1;
# CHECK: if (x>5
# CHECK: int A_2 = A[0];
# CHECK: for (int x
# CHECK: A_2 = A_2 + 1;
# CHECK: }
# CHECK: A[0] = A_2;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// references a Let var in a local scope which cannot be hoisted out of the
// loop.
TEST(Registerizer, RegisterizerLoopLetVar) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make({For::make(
x,
0,
10,
Block::make(
{Let::make(y, 30),
Store::make(a, {y}, Add::make(x, Load::make(a, {y}, 1)), 1)}))});
/*
* for (int x = 0; x < 10; x++) {
* int y = 30;
* A[y] = x + (A[y]);
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// references a Let var in an outer scope that does not prevent hoisting the
// initializer.
TEST(Registerizer, RegisterizerLoopLetVarOuter) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt =
Block::make({Let::make(y, 30),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {y}, Add::make(x, Load::make(a, {y}, 1)), 1)}))});
/*
* int y = 30;
* for (int x = 0; x < 10; x++) {
* A[y] = x + (A[y]);
* }
*/
stmt = registerize(stmt);
/*
* int y = 30;
* int A_1 = A[y];
* for (int x = 0; x < 10; x++) {
* A_1 = x + A_1;
* }
* A[y] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int y = 30;
# CHECK: int A_1 = A[y];
# CHECK: for (int x
# CHECK: A_1 = x + A_1;
# CHECK: A[y] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Okay so the registerizer generally goes after index flattening, but just in
// case. Test multi index registerization.
TEST(Registerizer, RegisterizerMultiDim) {
KernelScope kernel_scope;
BufHandle a("A", {3, 4, 5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0, 1, 2}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0, 1, 2}, Add::make(Load::make(a, {0, 1, 2}, 1), x), 1)}))});
/*
* A[0, 1, 2] = 0;
* for (int x = 0; x < 10; x++) {
* A[0, 1, 2] = (A[0, 1, 2]) + x;
* }
*/
stmt = registerize(stmt);
/*
* int A_1 = 0;
* for (int x = 0; x < 10; x++) {
* A_1 = x + A_1;
* }
* A[0, 1, 2] = A_1;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: int A_1 = 0;
# CHECK: for (int x = 0; x < 10; x++)
# CHECK-NOT: A[
# CHECK: A_1 =
# CHECK: A[0, 1, 2] = A_1;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// Wont registerize if only some dims match, but will still registerize distinct
// elements.
TEST(Registerizer, RegisterizerMultiDimPartial) {
KernelScope kernel_scope;
BufHandle a("A", {3, 4, 5}, kInt);
VarHandle x("x", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0, 1, 2}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0, 2, 2}, Add::make(Load::make(a, {0, 1, 4}, 1), x), 1)}))});
/*
* A[0, 1, 2] = 0;
* for (int x = 0; x < 10; x++) {
* A[0, 2, 2] = (A[0, 1, 4]) + x;
* }
*/
stmt = registerize(stmt);
/*
* A[0, 1, 2] = 0;
* int A_1 = A[0, 1, 4];
* int A_2 = A[0, 2, 2];
* for (int x = 0; x < 10; x++) {
* A_2 = x + A_1;
* }
* A[0, 2, 2] = A_2;
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: A[0, 1, 2] = 0;
# CHECK: int A_1 = A[0, 1, 4];
# CHECK: int A_2 = A[0, 2, 2];
# CHECK: for (
# CHECK: A_2 = x + A_1;
# CHECK: A[0, 2, 2] = A_2;)IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// If they could overlap across all dimensions we cannot registerize.
TEST(Registerizer, RegisterizerMultiDimOverlap) {
KernelScope kernel_scope;
BufHandle a("A", {3, 4, 5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0, 1, 2}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0, x, 2}, Add::make(Load::make(a, {y, 2, 2}, 1), x), 1)}))});
/*
* A[0, 1, 2] = 0;
* for (int x = 0; x < 10; x++) {
* A[0, x, 2] = (A[y, 2, 2]) + x;
* }
*/
std::ostringstream before;
before << *stmt;
// No change.
stmt = registerize(stmt);
std::ostringstream after;
after << *stmt;
ASSERT_EQ(before.str(), after.str());
}
// But, if one dimension is known to be distinct they do not overlap.
TEST(Registerizer, RegisterizerMultiDimPartialOverlap) {
KernelScope kernel_scope;
BufHandle a("A", {3, 4, 5}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
Stmt* stmt = Block::make(
{Store::make(a, {0, 1, 2}, 0, 1),
For::make(
x,
0,
10,
Block::make({Store::make(
a, {0, x, 2}, Add::make(Load::make(a, {y, 2, 4}, 1), x), 1)}))});
/*
* A[0, 1, 2] = 0; <---- 2nd dim overlaps with store.
* for (int x = 0; x < 10; x++) {
* A[0, x, 2] = (A[y, 2, 4]) + x; <---- 3rd dim has constant diff.
* }
*/
stmt = registerize(stmt);
/*
* A[0, 1, 2] = 0;
* int A_1 = A[y, 2, 4];
* for (int x = 0; x < 10; x++) {
* A[0, x, 2] = A_1 + x;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: A[0, 1, 2] = 0;
# CHECK: int A_1 = A[y, 2, 4];
# CHECK: for (
# CHECK: A[0, x, 2] = A_1 + x;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// A 3D reduction with different input dimensionality.
TEST(Registerizer, RegisterizerMultiDim3DReduction1) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10, 10}, kInt);
BufHandle c("C", {10, 10, 10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
VarHandle z("z", kInt);
Stmt* stmt = For::make(
x,
0,
10,
For::make(
y,
0,
10,
For::make(
z,
0,
10,
Store::make(
c,
{x, y, z},
Add::make(
Load::make(c, {x, y, z}, 1),
Mul::make(
Load::make(b, {x, y}, 1), Load::make(a, {x}, 1))),
1))));
/*
* for (int x = 0; x < 10; x++) {
* for (int y = 0; y < 10; y++) {
* for (int z = 0; z < 10; z++) {
* C[x, y, z] = (C[x, y, z]) + (B[x, y]) * (A[x]);
* }
* }
* }
*/
// We can registerize the A and B access since they can be hoisted before
// hitting a dependent loop var.
stmt = registerize(stmt);
/*
* for (int x = 0; x < 10; x++) {
* int A_1 = A[x];
* for (int y = 0; y < 10; y++) {
* int B_1 = B[x, y];
* for (int z = 0; z < 10; z++) {
* C[x, y, z] = A_1 * B_1 + (C[x, y, z]);
* }
* }
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: for (int x
# CHECK: int A_1 = A[x];
# CHECK: for (int y
# CHECK: int B_1 = B[x, y];
# CHECK: for (int z
# CHECK: C[x, y, z] = A_1 * B_1 + (C[x, y, z]);
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
// A 3D reduction with the same smaller dimensionality using different loop
// vars.
TEST(Registerizer, RegisterizerMultiDim3DReduction2) {
KernelScope kernel_scope;
BufHandle a("A", {10}, kInt);
BufHandle b("B", {10}, kInt);
BufHandle c("C", {10}, kInt);
VarHandle x("x", kInt);
VarHandle y("y", kInt);
VarHandle z("z", kInt);
Stmt* stmt = For::make(
x,
0,
10,
For::make(
y,
0,
10,
For::make(
z,
0,
10,
Store::make(
c,
{x},
Add::make(
Load::make(c, {x}, 1),
Mul::make(Load::make(b, {y}, 1), Load::make(a, {x}, 1))),
1))));
/*
* for (int x = 0; x < 10; x++) {
* for (int y = 0; y < 10; y++) {
* for (int z = 0; z < 10; z++) {
* C[x] = (C[x]) + (B[y]) * (A[x]);
* }
* }
* }
*/
// We can registerize all accesses, the A and C access can be hoisted to the
// outer loop since they depend only on it's loop var while the B can only be
// raised to the loop of y.
stmt = registerize(stmt);
/*
* for (int x = 0; x < 10; x++) {
* int A_1 = A[x];
* int C_1 = C[x];
* for (int y = 0; y < 10; y++) {
* int B_1 = B[y];
* for (int z = 0; z < 10; z++) {
* C_1 = B_1 * A_1 + C_1;
* }
* }
* C[x] = C_1;
* }
*/
std::ostringstream oss;
oss << *stmt;
const std::string& verification_pattern =
R"IR(
# CHECK: for (int x
# CHECK: int A_1 = A[x];
# CHECK: int C_1 = C[x];
# CHECK: for (int y
# CHECK: int B_1 = B[y];
# CHECK: for (int z
# CHECK: C_1 = B_1 * A_1 + C_1;
# CHECK: }
# CHECK: }
# CHECK: C[x] = C_1;
# CHECK: })IR";
torch::jit::testing::FileCheck().run(verification_pattern, oss.str());
}
} // namespace jit
} // namespace torch