End-to-end BERT GEMM lowering fails due to layout/packing mismatch in CKKS pipeline

Hi, I’m new to HEIR and am trying to run an end-to-end flow on a (simplified) BERT layer. My current pipeline is: 1. Use `torch-mlir` to lower a PyTorch BERT layer to the linalg dialect 2. Run `heir-opt` with `--torch-linalg-to-ckks` (backend: `openfhe`) 3. Use `heir-translate` to generate code As a starting point, I reduced the BERT layer to individual kernels (e.g., GEMM, GELU). However, I’m encountering multiple errors when processing GEMM variants, and I suspect they may be related to the packing/layout system. **Reproducer (GEMM variant)** ```milr #map = affine_map<(d0, d1, d2) -> (d1, d2)> #map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> module { func.func @forward(%arg0: tensor<1x128x768xf32> {secret.secret}) -> tensor<1x128x64xf32> { %cst = arith.constant 0.000000e+00 : f32 %cst_0 = arith.constant dense_resource<torch_tensor_64_768_torch.float32> : tensor<64x768xf32> %0 = tensor.empty() : tensor<768x64xf32> %transposed = linalg.transpose ins(%cst_0 : tensor<64x768xf32>) outs(%0 : tensor<768x64xf32>) permutation = [1, 0] %1 = tensor.empty() : tensor<1x768x64xf32> %2 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%transposed : tensor<768x64xf32>) outs(%1 : tensor<1x768x64xf32>) { ^bb0(%in: f32, %out: f32): linalg.yield %in : f32 } -> tensor<1x768x64xf32> %3 = tensor.empty() : tensor<1x128x64xf32> %4 = linalg.fill ins(%cst : f32) outs(%3 : tensor<1x128x64xf32>) -> tensor<1x128x64xf32> %5 = linalg.batch_matmul ins(%arg0, %2 : tensor<1x128x768xf32>, tensor<1x768x64xf32>) outs(%4 : tensor<1x128x64xf32>) -> tensor<1x128x64xf32> return %5 : tensor<1x128x64xf32> } } ``` **Observed issues** - The constant weight matrix is initialized with row-major dense packing. - The MatmulDiagonal kernel appears to expect a per-row layout. - As a result, a convert_layout op is inserted. However, this leads to a runtime failure during the ConvertToCiphertextSemantics pass: `slice along dimension 0 runs out-of-bounds` I also experimented with other GEMM variants and encountered additional layout-related errors, such as mismatches between the layout of a linalg.fill output and the expected layout for linalg.matmul. **Questions** - What is the current status of the packing/layout system? - Is this kind of layout mismatch expected at the moment, or am I missing a required preprocessing step? - Are there working examples of GEMM (or BERT-like workloads) that go through this pipeline?