def _pad(*indices):
not_zero = []
index_tuple = []
for i in range(n):
if pad_before[i] == 0 and pad_after[i] == 0:
index_tuple.append(indices[i])
else:
index_tuple.append(indices[i] - pad_before[i])
not_zero.append(indices[i] >= pad_before[i])
not_zero.append(indices[i] < data.shape[i] + pad_before[i])
if not_zero:
not_zero = tvm.all(*not_zero)
return tvm.select(not_zero, data[tuple(index_tuple)], pad_value)
return data[tuple(index_tuple)]
def kernel(A):
return hcl.compute((1, 1, 10, 10), lambda i, c, x, y: hcl.select(
tvm.all(x < 8, y < 8), A[i, c, x, y], 0), "B", dtype)