Incorrect Results for Conv2d

[CStarn]

I have been trying to compile and run the following minimal linalg.conv_2d example:

bazel run //tools:heir-opt -- --annotate-module="backend=openfhe scheme=ckks" --torch-linalg-to-ckks=ciphertext-degree=1024 --scheme-to-openfhe 
bazel run //tools:heir-translate -- --emit-openfhe-pke-header --openfhe-include-type=source-relative 
bazel run //tools:heir-translate -- --emit-openfhe-pke --openfhe-include-type=source-relative 

module {
  func.func @convolution(
    %arg0: tensor<3x3xf32> {secret.secret},
    %filter1: tensor<2x2xf32>
  ) -> tensor<2x2xf32> {
  
    %output = tensor.empty() : tensor<2x2xf32>
    %result = linalg.conv_2d 
      ins(%arg0, %filter1 : tensor<3x3xf32>, tensor<2x2xf32>)
      outs(%output : tensor<2x2xf32>) -> tensor<2x2xf32>
    
    return %result : tensor<2x2xf32>
  }
}

To be able to compile this example I had to implement the arith.floordivsi operation in the OpenFHEEmitter: https://github.com/Fraunhofer-AISEC/heir/commit/20360dee8dcea8d98625558590d4dd8e6344c5f1.

Running the generated code with the following snippet works but yields wrong outputs as listed below.

#include "test_conv2d_2x2.h"
#include <iostream>
#include <vector>
#include <iomanip>

int main() {
    CryptoContextT cc = convolution__generate_crypto_context();
    
    auto keyPair = cc->KeyGen();
    PublicKeyT publicKey = keyPair.publicKey;
    PrivateKeyT privateKey = keyPair.secretKey;
    
    cc = convolution__configure_crypto_context(cc, privateKey);
    
    std::vector<float> input_data(9, 0.1f);
    
    std::vector<float> filter_weights(4, 0.1f);
    
    std::cout << "Input (3x3):" << std::endl;
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            std::cout << std::fixed << std::setprecision(1) 
                      << input_data[i * 3 + j] << " ";
        }
        std::cout << std::endl;
    }
    std::cout << std::endl;
    
    std::cout << "Filter (2x2):" << std::endl;
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < 2; j++) {
            std::cout << std::fixed << std::setprecision(1) 
                      << filter_weights[i * 2 + j] << " ";
        }
        std::cout << std::endl;
    }
    std::cout << std::endl;
    
    std::vector<MutableCiphertextT> encrypted_input =
            convolution__encrypt__arg0(cc, input_data, publicKey);
    
    std::vector<MutableCiphertextT> encrypted_result =
            convolution(cc, encrypted_input, filter_weights);
    
    std::vector<float> decrypted_result =
            convolution__decrypt__result0(cc, encrypted_result, privateKey);
    
    std::cout << "Output (2x2):" << std::endl;
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < 2; j++) {
            std::cout << std::fixed << std::setprecision(6) 
                      << decrypted_result[i * 2 + j] << " ";
        }
        std::cout << std::endl;
    }
    
    return 0;
}

I would expect that all values in the result should be equal, but they aren't.

Result:

Input (3x3):
0.1 0.1 0.1 
0.1 0.1 0.1 
0.1 0.1 0.1 

Filter (2x2):
0.1 0.1 
0.1 0.1 

Output (2x2):
0.040000 0.040000 
0.030000 0.030000 

@asraa It would be great if you could have a look at this.

[asraa]

Thanks for the test case! I'll debug it today. It might be something to do with padding - the example https://github.com/google/heir/tree/aae0d07bb44a43e660e72c5a6567b1a8e70cfc83/tests/Examples/openfhe/ckks/conv_2dis 4x4 with a 3x3 kernel and seemed ok, so maybe there's something going on with padding a 3x3 data element, or a 2x2 filter.

[asraa]

Hey! Thanks so much for posting this issue, it was indeed a bug - we treated the indexes to access the tensors in the IR as unsigned integers (index types in MLIR) whereas the ISL library that we use to generate the layout assignments uses signed integer types. Updating our ISL to MLIR conversion to keep the signed integer types fixed the problem :)

This PR has the fix and also includes that floordivsi patch you made https://github.com/google/heir/pull/2433