Questions on C++ extension for customized convolution layers

Hi, I’m working on creating a C++ extension for a customized depthwise convolution layer on CPU. Besides all the parameters used in at::Tensor _convolution() in pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp, I have another input which is a mask in the same shape of input channel (const Tensor& mask). I met couple of problems when implementing this new convolution extension:

  1. Although I have used namespace at and native as used in pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp, I’m still facing the problem : error: no type named ‘IntArrayRef’ in namespace ‘at’ May I ask where is this IntArrayRef defined?
  2. I was not be able to use the dispatch as what the source code did at pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp line 21: DEFINE_DISPATCH(convolution_depthwise3x3_winograd_stub);
    The reason I’m using it is because I also want to call convolution_depthwise3x3_winograd_stub();
  3. This may sounds like a silly one, my convolution is executed based on the zeros and ones in my mask, so I was trying to use if-statements. It seems I was not be able to do some comparisons like:
    if (mask[i] == 0) or if (mask[i] == torch::zeros(shape)), it will return an error of error: value of type ‘at::Tensor’ is not contextually convertible to ‘bool’. May I ask how should I do it?

So the first question, I was able to solve it. According to, IntArrayRef is declared under name space of c10. When I change it to ArrayRef<int64_t>, it seems work.

The new question is that, in this lane,
a function expand_param_if_needed() was used but not declared ahead of time. Any clues on where it comes from?

expand_param_if_needed() is in aten/src/ATen/native/utils/ParamUtils.h

The third question, should be using mask[i].item()

It seems you have figured out most of your questions. Regarding DEFINE_DISPATCH, it seems that it is defined here:

#define DEFINE_DISPATCH(name) struct name name

If you get ATen related header included, you should be able to use it, right?

Hi, yeah. Thank you for the help. I currently solved it!

Hello, I am trying to implement something similar. It would be a great help if you could provide your code snippet here. Thanks in advance.