# CHECK-NEXT: B[)IR");
}
-TEST(LoopNest, ScheduleInlineBufferIndicesWithCast) {
+TEST(LoopNest, ScheduleInlineConsumerIndicesWithCast) {
// Input IR:
// for (int64_t i = 0; i < 100; i++) {
// A[0ll,i] = i * 500ll;
auto par = Block::make({forI, forJ});
LoopNest l(par, {b_buf.node()});
- l.computeInline(a_buf.node());
+ ASSERT_TRUE(l.computeInline(a_buf.node()));
+
+ checkIR(l.root_stmt(), R"IR(
+ # CHECK: for (int64_t j = 0; j < 100; j++) {
+ # CHECK: B[0ll, j] = j * 500ll + j * 100ll;
+ # CHECK: })IR");
+}
+
+TEST(LoopNest, ScheduleInlineProducerIndicesWithCast) {
+ // Input IR:
+ // for (int64_t i = 0; i < 100; i++) {
+ // A[(int64_t)0,i] = i * 500ll;
+ // }
+ // for (int64_t j = 0; j < 100; j++) {
+ // B[0ll,j] = A[0ll, j] + j * 100ll;
+ // }
+ BufHandle a_buf("A", {20, 100}, kLong);
+ BufHandle b_buf("B", {20, 100}, kLong);
+ VarHandle i("i", kLong);
+ VarHandle j("j", kLong);
+ auto forI = For::make(
+ i,
+ 0,
+ 100,
+ Store::make(a_buf, {0, i}, Mul::make(i, static_cast<int64_t>(500))));
+ auto forJ = For::make(
+ j,
+ 0,
+ 100,
+ Store::make(
+ b_buf,
+ {static_cast<int64_t>(0), j},
+ Add::make(
+ Load::make(a_buf, {static_cast<int64_t>(0), j}),
+ Mul::make(j, static_cast<int64_t>(100)))));
+ auto par = Block::make({forI, forJ});
+
+ LoopNest l(par, {b_buf.node()});
+ ASSERT_TRUE(l.computeInline(a_buf.node()));
checkIR(l.root_stmt(), R"IR(
# CHECK: for (int64_t j = 0; j < 100; j++) {
if (auto index_var = to<Var>(i)) {
index_vars_.insert(index_var);
producer_index_vars_.push_back(index_var);
- } else if (intValue(i)) {
+ } else {
// If the index can be a constant, then that dimension must have size 1
// (since we don't support in-place writes). Resolves issue 52581.
- if (*intValue(i) != 0) {
+ auto index_val = evalInt(i);
+ if (!index_val || *index_val != 0) {
success_ = false;
+ break;
}
producer_index_vars_.push_back(nullptr);
- } else {
- // Cannot inline Buf with compound indices
- success_ = false;
}
}
}