add conv2d conv1d forward function #166
Conversation
|
There are new conflicts needed to be resolve. Please see to that. |
src/flag_gems/ops/conv1d.py
Outdated
| + (tl.arange(0, BLOCK_C_IN) * padded_width_input)[:, None] | ||
| + i * stride_width | ||
| + tl.arange(0, BLOCK_W) | ||
| ) |
There was a problem hiding this comment.
doesn't support width_kernel that is not power of 2?
There was a problem hiding this comment.
As Continuous data loading problem. Currently, block only supports powers of 2, so the h w of the kernel only supports this specification
src/flag_gems/ops/conv1d.py
Outdated
| weight_value = tl.load(weight + weight_offset, mask=mask_weight, other=0) | ||
| input_value = tl.reshape(input_value, (BLOCK_N, BLOCK_OUT_WEIGHT)) | ||
| weight_value = tl.reshape(weight_value, (BLOCK_OUT_WEIGHT, BLOCK_O)) | ||
| accumulator = tl.dot(input_value, weight_value, allow_tf32=False) |
There was a problem hiding this comment.
accumulator redefined
tests/test_reduction_ops.py
Outdated
| @pytest.mark.parametrize("kernel", [(17, 2, 2)]) | ||
| @pytest.mark.parametrize("stride", [2]) | ||
| @pytest.mark.parametrize("padding", [1]) | ||
| @pytest.mark.parametrize("dtype", [torch.float32]) |
src/flag_gems/ops/conv2d.py
Outdated
| ) | ||
| input_ci_offset = offset_ci[:, None, None] * height_input * width_input | ||
| input_group_offset = ( | ||
| pid_group[None, None, None, None] * c_input * width_input * height_input |
There was a problem hiding this comment.
initializing input_group_offset as a scalar is okay
src/flag_gems/ops/conv2d.py
Outdated
| input_col = mm(out_grad, weight_reshape.T) | ||
|
|
||
| # return dx,None,None,None,None,None,None | ||
| conv2d_col2img[grid]( |
There was a problem hiding this comment.
maybe we could fuse mm and col2img together as bwd kernel
FatJhon
force-pushed
the
dev_xcoresigma_jiangbin_conv
branch
from
121285e
to
5adc6e6
Compare
| @@ -16,6 +16,9 @@ | |||
| from .bmm import bmm | |||
| from .cat import cat | |||
| from .clamp import clamp, clamp_tensor | |||
| from .conv1d import conv1d | |||
| from .conv2d import conv2d | |||
| from .conv_depthwise2d import _conv_depthwise2d | |||
There was a problem hiding this comment.
Is conv_depthwise2d implemented?
| padding_width = padding[0] | ||
| else: | ||
| padding_width = padding | ||
| return conv2d( |
There was a problem hiding this comment.
I suggest implementing a kernel function for conv1d specifically. Calling conv2d will cost additional runtime.
src/flag_gems/ops/conv2d.py
Outdated
| accum = tl.zeros((BLOCK_NI_HO_WO, BLOCK_CO), dtype=tl.float32) | ||
|
|
||
| for h in range(kernel_height): | ||
| for w in range(kernel_width): |
There was a problem hiding this comment.
since kernel_height and kernel_width are processed as loop iterator range, why not support those non-power-of-2?
|
|
||
| class Conv2d(torch.autograd.Function): | ||
| @staticmethod | ||
| def forward(ctx, input, weight, bias, stride, padding, dilation, groups): |
There was a problem hiding this comment.
Dilation is not used at all, so it does not support dilation other than 1?
There was a problem hiding this comment.
Dilation is not used.
There was a problem hiding this comment.
I think we should raise an error if dilation > 1.
There was a problem hiding this comment.
added dilation function
src/flag_gems/ops/conv2d.py
Outdated
| Returns: | ||
| Output size of 2D convolution. | ||
| """ | ||
| return (in_size + 2 * padding - kernel_size) // stride + 1 |
There was a problem hiding this comment.
Does it support asymmetric padding? I suppose not. Alright, torch's convolution does not support asymmetric padding.
|
|
||
| return ( | ||
| input, | ||
| None, |
There was a problem hiding this comment.
Gradient of weight and bias should also be computed.
| # default conv shape for input and weight stride padding groups | ||
| # default Ni Ci Hi WI Co Hk Wk stride padding groups | ||
| ConvBenchmark: | ||
| shapes: |
There was a problem hiding this comment.
5 shapes are enough for core mode.
We have completed the forward fuction of the conv1d, conv2d conv2d_depthwise operators.