vectorize schedule didn't generate vectorize instructions for CUDA

[jinderek]

I use the vectorize schedule, but I look into PTX generated, there are no vectorize instructions(ld.global.v4 or st.global.v4 ....).
Did I miss something?

.cpp code

int main(int argc, char **argv) {
  Var x("x"), y("y");
  Func test("test_in_vectors");
  test(x, y) = x + y;

  Var x_outer, x_inner;
  test.split(x, x_outer, x_inner, 4);
  test.vectorize(x_inner);

  Target target = Halide::get_host_target();
  target.set_feature(Halide::Target::CUDA);
  test.gpu_blocks(y).gpu_threads(x_outer);

  Buffer<int> output = test.realize(8, 8, target);

}

The PTX generated:

// %bb.0:                               // %entry
        ld.param.u32    %r29, [kernel_test_in_vectors_s0_y___block_id_x_param_6];
        ld.param.u32    %r28, [kernel_test_in_vectors_s0_y___block_id_x_param_5];
        ld.param.u32    %r30, [kernel_test_in_vectors_s0_y___block_id_x_param_0];
        ld.param.u64    %rd2, [kernel_test_in_vectors_s0_y___block_id_x_param_7];
        cvta.to.global.u64      %rd1, %rd2;
        mov.u32         %r1, %ctaid.x;
        mov.u32         %r2, %tid.x;
        setp.ge.s32     %p1, %r2, %r30;
        @%p1 bra        LBB0_2;
// %bb.1:                               // %true_bb
        ld.param.u32    %r27, [kernel_test_in_vectors_s0_y___block_id_x_param_4];
        ld.param.u32    %r24, [kernel_test_in_vectors_s0_y___block_id_x_param_1];
        shl.b32         %r32, %r2, 2;
        add.s32         %r33, %r1, %r24;
        add.s32         %r37, %r33, %r32;
        add.s32         %r40, %r37, 3;
        add.s32         %r39, %r37, 2;
        add.s32         %r38, %r37, 1;
        add.s32         %r34, %r1, %r28;
        mad.lo.s32      %r35, %r34, %r29, %r27;
        add.s32         %r36, %r35, %r32;
        bra.uni         LBB0_3;
LBB0_2:                                 // %false_bb
        ld.param.u32    %r26, [kernel_test_in_vectors_s0_y___block_id_x_param_3];
        ld.param.u32    %r25, [kernel_test_in_vectors_s0_y___block_id_x_param_2];
        add.s32         %r37, %r1, %r25;
        add.s32         %r40, %r37, 3;
        add.s32         %r39, %r37, 2;
        add.s32         %r38, %r37, 1;
        add.s32         %r31, %r1, %r28;
        mad.lo.s32      %r36, %r31, %r29, %r26;
LBB0_3:                                 // %after_bb
        mul.wide.s32    %rd3, %r36, 4;
        add.s64         %rd4, %rd1, %rd3;
        st.global.u32   [%rd4+4], %r38;
        st.global.u32   [%rd4], %r37;
        st.global.u32   [%rd4+8], %r39;
        st.global.u32   [%rd4+12], %r40;

HTML (But we can see ramp here)

    gpu_block<CUDA> (test_in_vectors.s0.y.__block_id_x, 0, test_in_vectors.extent.1) {
      gpu_thread<CUDA> (.__thread_id_x, 0, t13) {
        if ((.__thread_id_x < t14)) {
          test_in_vectors[ramp(((.__thread_id_x*4) + (((test_in_vectors.min.1 + test_in_vectors.s0.y.__block_id_x)*test_in_vectors.stride.1) + t7)), 1, 4)] = ramp(((.__thread_id_x*4) + (t15 + test_in_vectors.s0.y.__block_id_x)), 1, 4)
        } else {
          test_in_vectors[ramp((((test_in_vectors.min.1 + test_in_vectors.s0.y.__block_id_x)*test_in_vectors.stride.1) + t17), 1, 4)] = ramp((t16 + test_in_vectors.s0.y.__block_id_x), 1, 4)
        }
      }
    }
    _halide_buffer_set_device_dirty(test_in_vectors.buffer, (uint1)1)
  }
}

[abadams]

CUDA has alignment requirements to actually generate vector ops unfortunately, so in the IR you're looking for the aligned tag. Here's one way to give halide+llvm sufficient information such that the stores should be known to be aligned:

  Var x_outer, x_inner;
  test.split(x, x_outer, x_inner, 4, TailStrategy::RoundUp);
  test.vectorize(x_inner);
  test.output_buffer().dim(0).set_min(0);

[jinderek]

Thanks for the detailed reply.
After changing to float and using compute_root, I can see vectorize load(ld.global.nc.v2.f32) in the ptx, but the store is still not vectorize store.
How to schedule to generate code with vectorize store?

.cpp

int main(int argc, char **argv) {
  Func a;
  Func test;
  Var x, y;
  a(x) = cast<float>(x);
  test(x, y) = Halide::max(a(x), 0);

  a.compute_root();

  Var x_outer, x_inner;
  test.split(x, x_outer, x_inner, 2, TailStrategy::RoundUp);
  test.vectorize(x_inner);

  Target target = Halide::get_host_target();
  target.set_feature(Halide::Target::CUDA);
//  target.set_feature(Halide::Target::CUDACapability50);
  test.gpu_blocks(y).gpu_threads(x_outer);

  int size = 1024;
  int x_size = 1024;
  test.bound(x, 0, x_size).bound(y, 0, size);
  test.output_buffer().dim(0).set_min(0);
  Buffer<float> output = test.realize(x_size, size, target);
}

PTX code

// %bb.0:                               // %entry
        ld.param.u32    %r1, [kernel_test_s0_y___block_id_x_param_0];
        ld.param.u64    %rd1, [kernel_test_s0_y___block_id_x_param_2];
        cvta.to.global.u64      %rd2, %rd1;
        ld.param.u64    %rd3, [kernel_test_s0_y___block_id_x_param_1];
        cvta.to.global.u64      %rd4, %rd3;
        mov.u32         %r2, %ctaid.x;
        mov.u32         %r3, %tid.x;
        shl.b32         %r4, %r3, 1;
        mul.wide.u32    %rd5, %r4, 4;
        add.s64         %rd6, %rd4, %rd5;
        ld.global.v2.f32        {%f1, %f2}, [%rd6];
        max.ftz.f32     %f3, %f1, 0f00000000;
        max.ftz.f32     %f4, %f2, 0f00000000;
        mad.lo.s32      %r5, %r2, %r1, %r4;
        mul.wide.s32    %rd7, %r5, 4;
        add.s64         %rd8, %rd2, %rd7;
        st.global.f32   [%rd8+4], %f4;
        st.global.f32   [%rd8], %f3;
        ret;

CUDA has alignment requirements to actually generate vector ops unfortunately, so in the IR you're looking for the aligned tag. Here's one way to give halide+llvm sufficient information such that the stores should be known to be aligned:

  Var x_outer, x_inner;
  test.split(x, x_outer, x_inner, 4, TailStrategy::RoundUp);
  test.vectorize(x_inner);
  test.output_buffer().dim(0).set_min(0);