The schedule of two unrelated reduce sum

[Vernlium]

Hi all,
I have a case like this:

class TestGenerator : public Halide::Generator<TestGenerator> {
public:
    Input<Buffer<int16_t>> input{"input", 2};
    Output<int16_t> totalSum{"totalSum"};

    void generate() {
        Var x{"x"}, y{"y"};
        Func input_c{"input_c"};
        Func add_1{"add_1"};
        Func mul_2{"mul_2"};
        Func sum_1{"sum_1"};
        Func sum_2{"sum_2"};
        Func totalSum_c{"totalSum_c"};
        RDom r(input);

        // algorithm
        input_c(x, y) = input(x, y);

        add_1(x, y) = input_c(x, y) + 1;

        mul_2(x, y) = input_c(x, y) * 2;

        sum_1() = cast<int16_t>(0);
        sum_2() = cast<int16_t>(0);

        sum_1() += add_1(r.x, r.y);
        sum_2() += mul_2(r.x, r.y);

        totalSum_c() = sum_1() + sum_2();
        totalSum() = totalSum_c();

        input.dim(0).set_bounds(0, 32);
        input.dim(1).set_bounds(0, 64);

        // CPU schedule.
        int h_factor = 8;
        int w_factor = 8;

        RVar rxOuter("rxOuter"), rxInner("rxInner");
        RVar ryOuter("ryOuter"), ryInner("ryInner");

        RVar r_sum_x(sum_1.update(0).get_schedule().dims()[0].var);
        RVar r_sum_y(sum_1.update(0).get_schedule().dims()[1].var);

        std::cout << "r_sum_x.name=" << r_sum_x.name() << ", r_sum_y.name=" << r_sum_x.name() << std::endl;

        sum_1.update(0).tile(r_sum_x, r_sum_y, rxOuter, ryOuter, rxInner, ryInner, w_factor, h_factor);

        RVar r_sum_x_2(sum_2.update(0).get_schedule().dims()[0].var);
        RVar r_sum_y_2(sum_2.update(0).get_schedule().dims()[1].var);

        sum_2.update(0).tile(r_sum_x_2, r_sum_y_2, rxOuter, ryOuter, rxInner, ryInner, w_factor, h_factor);

        add_1.compute_at(sum_1, rxOuter);
        sum_1.update(0)
            .atomic();

        input_c.compute_at(sum_2, rxOuter);
        mul_2.compute_at(sum_2, rxOuter);
        sum_2.update(0)
            .atomic();

        sum_1.update(0).compute_with(sum_2.update(0), rxOuter);
        sum_1.compute_root();
        sum_2.compute_root();
        totalSum_c.compute_root();
        totalSum.compute_root();
    }
};

HALIDE_REGISTER_GENERATOR(TestGenerator, test_generator)

I use the compute_with merge the forloop of sum_1 and sum_2.

I get the following IR(the IR is after Unrolling):

 allocate sum_2[int16]
 allocate sum_1[int16]
 produce sum_1 {
  produce sum_2 {
   sum_1[0] = (int16)0
   sum_2[0] = (int16)0
   for (sum_2.s1.r6$y.ryOuter, 0, 8) {
    for (sum_2.s1.r6$x.rxOuter, 0, 4) {
     allocate input_c[int16 * 32 * 64] if (uint1)likely((uint1)1)
     produce input_c {
      for (input_c.s0.y, 0, 64) {
       for (input_c.s0.x, 0, 32) {
        input_c[(input_c.s0.y*32) + input_c.s0.x] = input[(input.stride.1*input_c.s0.y) + input_c.s0.x]
       }
      }
     }
     allocate mul_2[int16 * 8 * 8] if (uint1)likely((uint1)1)
     produce mul_2 {
      consume input_c {
       for (mul_2.s0.y, sum_2.s1.r6$y.ryOuter*8, 8) {
        for (mul_2.s0.x, sum_2.s1.r6$x.rxOuter*8, 8) {
         mul_2[((mul_2.s0.y - (sum_2.s1.r6$y.ryOuter*8))*8) + (mul_2.s0.x - (sum_2.s1.r6$x.rxOuter*8))] = input_c[(mul_2.s0.y*32) + mul_2.s0.x]*(int16)2
        }
       }
      }
     }
     consume mul_2 {
      for (sum_2.s1.r6$y.ryInner, 0, 8) {
       for (sum_2.s1.r6$x.rxInner, 0, 8) {
        atomic {
          sum_2[0] = sum_2[0] + mul_2[(sum_2.s1.r6$y.ryInner*8) + sum_2.s1.r6$x.rxInner]
        }
       }
      }
     }
     allocate add_1[int16 * 8 * 8] if (uint1)likely((uint1)1)
     produce add_1 {
      consume input_c {
       for (add_1.s0.y, sum_2.s1.r6$y.ryOuter*8, 8) {
        for (add_1.s0.x, sum_2.s1.r6$x.rxOuter*8, 8) {
         add_1[((add_1.s0.y - (sum_2.s1.r6$y.ryOuter*8))*8) + (add_1.s0.x - (sum_2.s1.r6$x.rxOuter*8))] = input_c[(add_1.s0.y*32) + add_1.s0.x] + (int16)1
        }
       }
      }
     }
     consume add_1 {
      for (sum_1.s1.r6$y.ryInner, 0, 8) {
       for (sum_1.s1.r6$x.rxInner, 0, 8) {
        atomic {
          sum_1[0] = sum_1[0] + add_1[(sum_1.s1.r6$y.ryInner*8) + sum_1.s1.r6$x.rxInner]
        }
       }
      }
     }
    }
   }
  }
 }
 allocate totalSum_c[int16]
 produce totalSum_c {
  consume sum_1 {
   consume sum_2 {
    totalSum_c[0] = sum_1[0] + sum_2[0]
   }
  }
 }
 produce totalSum {
  consume totalSum_c {
   totalSum[0] = totalSum_c[0]
  }
 }
}

Pay attention to the produce input_c:

     allocate input_c[int16 * 32 * 64] if (uint1)likely((uint1)1)
     produce input_c {
      for (input_c.s0.y, 0, 64) {
       for (input_c.s0.x, 0, 32) {
        input_c[(input_c.s0.y*32) + input_c.s0.x] = input[(input.stride.1*input_c.s0.y) + input_c.s0.x]
       }
      }
     }

I expect the allocate size of input_c is 8 * 8, and the forloop extent of input_c are 8 and 8. But the allocate size of input_c is 32 * 64, that is the input_c need compute all data for every outer forloop.

So, is there a bug of Halide backend Or is there something wrong with my schedule?

--

I find a way to slove the problem.

clone_in the input_c.

The schedule is：

       // CPU schedule.
        int h_factor = 8;
        int w_factor = 8;

        RVar rxOuter("rxOuter"), rxInner("rxInner");
        RVar ryOuter("ryOuter"), ryInner("ryInner");

        RVar r_sum_x(sum_1.update(0).get_schedule().dims()[0].var);
        RVar r_sum_y(sum_1.update(0).get_schedule().dims()[1].var);

        std::cout << "r_sum_x.name=" << r_sum_x.name() << ", r_sum_y.name=" << r_sum_x.name() << std::endl;

        sum_1.update(0).tile(r_sum_x, r_sum_y, rxOuter, ryOuter, rxInner, ryInner, w_factor, h_factor);

        RVar r_sum_x_2(sum_2.update(0).get_schedule().dims()[0].var);
        RVar r_sum_y_2(sum_2.update(0).get_schedule().dims()[1].var);

        sum_2.update(0).tile(r_sum_x_2, r_sum_y_2, rxOuter, ryOuter, rxInner, ryInner, w_factor, h_factor);

        Func input_c_clone("input_c_clone");
        input_c_clone = input_c.clone_in(add_1);

        input_c_clone.compute_at(sum_1, rxOuter);
        add_1.compute_at(sum_1, rxOuter);
        sum_1.update(0)
            .atomic();

        input_c.compute_at(sum_2, rxOuter);
        mul_2.compute_at(sum_2, rxOuter);
        sum_2.update(0)
            .atomic();

        sum_1.update(0).compute_with(sum_2.update(0), rxOuter);
        sum_1.compute_root();
        sum_2.compute_root();
        totalSum_c.compute_root();
        totalSum.compute_root();

The IR is：

 allocate sum_2[int16]
 allocate sum_1[int16]
 produce sum_1 {
  produce sum_2 {
   sum_1[0] = (int16)0
   sum_2[0] = (int16)0
   for (sum_2.s1.input$y.ryOuter, 0, 8) {
    for (sum_2.s1.input$x.rxOuter, 0, 4) {
     allocate input_c[int16 * 8 * 8] if (uint1)likely((uint1)1)
     produce input_c {
      for (input_c.s0.y, sum_2.s1.input$y.ryOuter*8, 8) {
       for (input_c.s0.x, sum_2.s1.input$x.rxOuter*8, 8) {
        input_c[((input_c.s0.y - (sum_2.s1.input$y.ryOuter*8))*8) + (input_c.s0.x - (sum_2.s1.input$x.rxOuter*8))] = input[(input.stride.1*input_c.s0.y) + input_c.s0.x]
       }
      }
     }
     allocate mul_2[int16 * 8 * 8] if (uint1)likely((uint1)1)
     produce mul_2 {
      consume input_c {
       for (mul_2.s0.y, sum_2.s1.input$y.ryOuter*8, 8) {
        for (mul_2.s0.x, sum_2.s1.input$x.rxOuter*8, 8) {
         mul_2[((mul_2.s0.y - (sum_2.s1.input$y.ryOuter*8))*8) + (mul_2.s0.x - (sum_2.s1.input$x.rxOuter*8))] = input_c[((mul_2.s0.y - (sum_2.s1.input$y.ryOuter*8))*8) + (mul_2.s0.x - (sum_2.s1.input$x.rxOuter*8))]*(int16)2
        }
       }
      }
     }
     consume mul_2 {
      for (sum_2.s1.input$y.ryInner, 0, 8) {
       for (sum_2.s1.input$x.rxInner, 0, 8) {
        atomic {
          sum_2[0] = sum_2[0] + mul_2[(sum_2.s1.input$y.ryInner*8) + sum_2.s1.input$x.rxInner]
        }
       }
      }
     }
     allocate input_c_clone_in_add_1$0[int16 * 8 * 8] if (uint1)likely((uint1)1)
     produce input_c_clone_in_add_1$0 {
      for (input_c_clone_in_add_1$0.s0.y, sum_2.s1.input$y.ryOuter*8, 8) {
       for (input_c_clone_in_add_1$0.s0.x, sum_2.s1.input$x.rxOuter*8, 8) {
        input_c_clone_in_add_1$0[((input_c_clone_in_add_1$0.s0.y - (sum_2.s1.input$y.ryOuter*8))*8) + (input_c_clone_in_add_1$0.s0.x - (sum_2.s1.input$x.rxOuter*8))] = input[(input.stride.1*input_c_clone_in_add_1$0.s0.y) + input_c_clone_in_add_1$0.s0.x]
       }
      }
     }
     allocate add_1[int16 * 8 * 8] if (uint1)likely((uint1)1)
     produce add_1 {
      consume input_c_clone_in_add_1$0 {
       for (add_1.s0.y, sum_2.s1.input$y.ryOuter*8, 8) {
        for (add_1.s0.x, sum_2.s1.input$x.rxOuter*8, 8) {
         add_1[((add_1.s0.y - (sum_2.s1.input$y.ryOuter*8))*8) + (add_1.s0.x - (sum_2.s1.input$x.rxOuter*8))] = input_c_clone_in_add_1$0[((add_1.s0.y - (sum_2.s1.input$y.ryOuter*8))*8) + (add_1.s0.x - (sum_2.s1.input$x.rxOuter*8))] + (int16)1
        }
       }
      }
     }
     consume add_1 {
      for (sum_1.s1.input$y.ryInner, 0, 8) {
       for (sum_1.s1.input$x.rxInner, 0, 8) {
        atomic {
          sum_1[0] = sum_1[0] + add_1[(sum_1.s1.input$y.ryInner*8) + sum_1.s1.input$x.rxInner]
        }
       }
      }
     }
    }
   }
  }
 }
 allocate totalSum_c[int16]
 produce totalSum_c {
  consume sum_1 {
   consume sum_2 {
    totalSum_c[0] = sum_1[0] + sum_2[0]
   }
  }
 }
 produce totalSum {
  consume totalSum_c {
   totalSum[0] = totalSum_c[0]
  }
 }
}

I get the expected allocate size and forloop extent. But the input_c have double calculation. Is there any way to eliminate this redundant calculation?

[vksnk]

Thanks a lot for detailed description! It does look like a bug / missing optimization in compute_with, I'll create an issue to track it and copy all info from here.

[Vernlium]

Thank you for your reply. How soon will this bug be solved?

[vksnk]

#6078 should fix it.