Simplify the if statement in backprop depthwise convolution

[weitliu]

Hi everyone. I am working on the backward method for the depth wise convolution. The implementation I currently can think of has a lot of tvm.select. Is there any way we could simplify the code?

def trans(b, i, j, c):
    global Out_grad_cond
    Out_grad_cond = tvm.compute(
        (batch, in_h, in_w, out_c),
        lambda bo, io, jo, co: tvm.select(tvm.all(io >= tvm.select(0<(i - filter_h + pad_h + stride_h) / stride_h,(i-filter_h+pad_h+stride_h)/stride_h,tvm.const(0)),
                                                  io <  tvm.select(0<((i + pad_h) / stride_h)+1-out_h, tvm.const(out_h - 1), (i + pad_h) / stride_h),
                                                  jo >= tvm.select(0<(j - filter_w + pad_w + stride_w) / stride_w,(j-filter_w+pad_w+stride_w)/stride_w,tvm.const(0)),
                                                  jo <  tvm.select(0<((j + pad_w) / stride_w)+1-out_w, tvm.const(out_w - 1), (j + pad_w) / stride_w)),
                                                  Out_grad[b, i, j, c], tvm.const(0.0)))

    di = tvm.reduce_axis((0, out_h-1), name='di')
    dj = tvm.reduce_axis((0, out_w-1), name='dj')
    dc = tvm.reduce_axis((0, channel_multiplier), name='dc')

    return tvm.sum(Out_grad_cond[b, di, dj, c*channel_multiplier + dc] * Filter[i+pad_h-di*stride_h, j+pad_w-dj*stride_w, c, dc],axis=[di,dj,dc])

In_grad = tvm.compute(
         (batch, in_h, in_w, in_c),
         lambda b, i, j, c: trans(b,i,j,c),
         name='In_grad')

Thank you!

[tqchen]

I think if the shape are constant, you only need a pad step, which is like out_grad_cond, but only requires check for four boundary conditions. This is like the pad step you might have in the input padding, and can be done efficiently if they are cached to temp.

cc @Huyuwei

[Huyuwei]

Looks like Out_grad_cond is doing two things: dilation (for stride larger than 1), and padding. Maybe try declare it as a dilation stage + a padding stage?

[tqchen]

makes sense, I think we can have a seperate padding and dilation computation declaration in TOPI, contribution is welcomed

[weitliu]

Interesting! I will take a look at it and report back later. Thank you!