torch.aten.avg_pool3d to linalg

[wu-s-john]

Required in model VideoResNet_vaiq_int8

%1473 = torch.aten.avg_pool3d %1471, %1472, %60, %551, %false, %false, %none : !torch.vtensor<[1,512,14,14,14],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[1,512,1,1,1],f32>

Interested in implement the lowering for AvgPool3dOp from Torch to linalg.

Seems that there is no native 3D convolution method for doing a sum or max in linalg. I would be happy to implement.

#215

llvm/torch-mlir#3030

[wu-s-john]

This was actually quite involved to implement. Namely, we cannot apply pooling_ndhwc_sum generically to ConvertAtenAvgPoolOp in Pooling.cpp.

In ConvertAtenAvgPoolOp::matchAndRewrite, it would eventually call the generic createPoolingOp. It would expect that the pooling would happen for a 5 dimensional matrix on it's last 3 dimensions (3, 4, 5). However, it does the pooling on (2, 3, 4) which is not what we want.

That being said, the best approach I thought was just writing out the entire convolution using linalg::Generic.

I thought it would be a bit cumbersome to constantly trying to achieve computational correctness, so I wanted to solve an easier problem. So, I implemented, I implemented the lowering, torch.aten.max_pool3d , instead.

[vivekkhandelwal1]

This was actually quite involved to implement. Namely, we cannot apply pooling_ndhwc_sum generically to ConvertAtenAvgPoolOp in Pooling.cpp.

In ConvertAtenAvgPoolOp::matchAndRewrite, it would eventually call the generic createPoolingOp. It would expect that the pooling would happen for a 5 dimensional matrix on it's last 3 dimensions (3, 4, 5). However, it does the pooling on (2, 3, 4) which is not what we want.

That being said, the best approach I thought was just writing out the entire convolution using linalg::Generic.

I thought it would be a bit cumbersome to constantly trying to achieve computational correctness, so I wanted to solve an easier problem. So, I implemented, I implemented the lowering, torch.aten.max_pool3d , instead.

Hi @wu-s-john, are you still working on this op?

[AmosLewis]

llvm/torch-mlir#3030