Support 1D convolution in HEIR

[j2kun]

This ticket is to support linalg.conv_1d in the HEIR compiler. This came up because Marc from Belfort was trying to lower a torch model that used 1d convolution. Together we did a deep dive on the layout system and he will contribute the implementation.

The basic idea is to reduce conv1d to Halevi-Shoup matvec with sparse diagonals, similar to the conv2d work.

The outline of supporting it:

• Add a filter-to-undiagonalized-matrix layout relation in lib/Utils/Layout, and then compose it with the matrix diagonalization relation
• Add a new function in lib/Kernel/KernelImplementation.h that constructs the arithmetic dag for the conv1d kernel, and uses implementHaleviShoup as a subroutine. This can be used to run the resulting kernel on test data by instantiating the template type for the ArithmeticDag with a concrete type (std::vector) and using EvalVisitor in a unit test.
• Add a new RewritePattern in ConvertToCiphertextSemantics.cpp (similar to the one for conv2d) that lowers conv1d using the arithmetic dag, and guards this by the kernel attribute (KernelName::MatvecDiagonal).
• Add support for conv1d in layout-propagation (I think this should be easy because the ciphertext input does not change its layout after going through this kernel).
• [optional]: Implement LayoutConversionHoistableOpInterface for linalg.conv_1d, which allows layout conversions to be hoistable through this op in layout-optimization.

[j2kun]

CC @mdgrs @AlexanderViand

[j2kun]

I was looking at the linalg dialect today and noticed that padding is not an attribute on most (all?) of the convolution ops. Since we talked about padding in the draft of #2899, I am wondering how @AlexanderViand or @mdgrs saw padding materialize in the IR?

[mdgrs]

I also saw 0 instance of padding in the convolution operations for the linalg dialect. My superficial understanding is that one would use a PadOp to add actual padding to the tensor in one of the passes. See https://github.com/llvm/llvm-project/blob/b48aa05f399d16faf90475cfb8131ae8693d8f97/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp#L2090

The pytorch Conv1D (as well as Conv2d) both expose a padding argument. This is also used in Orion, so I guess this is where we picked it up from. I did not actually see a padding argument when lowering from pytorch directly.

Finally, the conv2D op in HEIR does have a padding argument. As far as I can tell, when we convert from linalg.conv2d using the struct ConvertLinalgConv2D, the padding argument is hardcoded to 0. The presence of this argument might be related to #2883 , so perhaps @asraa can shed some more light here.

[AlexanderViand]

@mdgrs What does the mlir-torch-generated linalg IR look like if you change the padding settings in pytorch? Does it emit a linalg.pad?

[mdgrs]

Here is what it looks like if I just pad a tensor using pytorch:

       # x of dim 1 x 128
        x = nn.functional.pad(x,(1,1))

    %9 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%8, %8 : tensor<1x128xf32>, tensor<1x128xf32>) outs(%1 : tensor<1x128xf32>) {
    ^bb0(%in: f32, %in_8: f32, %out: f32):
      %15 = arith.mulf %in, %in_8 : f32
      linalg.yield %15 : f32
    } -> tensor<1x128xf32>
    %padded = tensor.pad %9 low[0, 1] high[0, 1] {
    ^bb0(%arg1: index, %arg2: index):
      tensor.yield %cst : f32
    } : tensor<1x128xf32> to tensor<1x130xf32>
    %10 = tensor.empty() : tensor<130x10xf32>
    %transposed_7 = linalg.transpose ins(%cst_1 : tensor<10x130xf32>) outs(%10 : tensor<130x10xf32>) permutation = [1, 0] 
    %11 = tensor.empty() : tensor<1x10xf32>
    %12 = linalg.fill ins(%cst : f32) outs(%11 : tensor<1x10xf32>) -> tensor<1x10xf32>
    %13 = linalg.matmul ins(%padded, %transposed_7 : tensor<1x130xf32>, tensor<130x10xf32>) outs(%12 : tensor<1x10xf32>) -> tensor<1x10xf32>
    %14 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} ins(%13, %cst_0 : tensor<1x10xf32>, tensor<10xf32>) outs(%11 : tensor<1x10xf32>) {
    ^bb0(%in: f32, %in_8: f32, %out: f32):
      %15 = arith.addf %in, %in_8 : f32
      linalg.yield %15 : f32
    } -> tensor<1x10xf32>
    return %14 : tensor<1x10xf32>

Unfortunately I am encountering issues lowering a pytorch Conv1D from the torch IR to linalg using torch-mlir: conv is declared as illegal. I will investigate further as we certainly need to fix this.

[mdgrs]

I found a test in torch-mlir that uses padding, and this is the output:

pytorch

class Conv1dWithSamePaddingModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        torch.manual_seed(0)
        self.conv = torch.nn.Conv1d(2, 10, 3, bias=False, padding="same")
        self.train(False)

    @export
    @annotate_args(
        [
            None,
            ([-1, -1, -1], torch.float32, True),
        ]
    )
    def forward(self, x):
        return self.conv(x)

Linalg IR

LINALG Backend IR
module {
  func.func @Conv1dWithSamePaddingModule(%arg0: tensor<5x2x10xf32>) -> tensor<5x10x10xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %cst_0 = arith.constant dense_resource<torch_tensor_10_2_3_torch.float32> : tensor<10x2x3xf32>
    %padded = tensor.pad %arg0 low[0, 0, 1] high[0, 0, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index):
      tensor.yield %cst : f32
    } : tensor<5x2x10xf32> to tensor<5x2x12xf32>
    %0 = tensor.empty() : tensor<5x10x10xf32>
    %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x10x10xf32>) -> tensor<5x10x10xf32>
    %2 = linalg.conv_1d_ncw_fcw {dilations = dense<1> : vector<1xi64>, strides = dense<1> : vector<1xi64>} ins(%padded, %cst_0 : tensor<5x2x12xf32>, tensor<10x2x3xf32>) outs(%1 : tensor<5x10x10xf32>) -> tensor<5x10x10xf32>
    return %2 : tensor<5x10x10xf32>
  }
}

i.e. it uses tensor.pad

So it looks to me like the tensor gets replaced with a tensor with more entries (of value 0). For the relation we implemented for conv1D, the behavior is not quite the same

• If datasize=5, padding = 1, then the number of columns of the kernel matrix would be 5.
• If datasize=7 (with two extra 0 values), then the number of columns would be 7.

So the kernel matrices would be different. The product between the kernel matrix and the data would be the same because the 0's align correctly.