google / google/heir

tensor_ext.remap doesn't handle shape expansion properly

Open
#2,847 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

This is derived from #2844

The reproducer with `--convert-to-ciphertext-semantics=ciphertext-size=16384`
```
#layout = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = i0 and (-i1 + slot) mod 1024 = 0 and 0 <= i0 <= 127 and 0 <= i1 <= 767 and 0 <= slot <= 16383 }">
#layout1 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : i0 = 0 and (-768i1 - i2 + slot + 16384*floor((768i1 + i2)/16384)) mod 131072 = 0 and 0 <= i1 <= 127 and 0 <= i2 <= 767 and 0 <= ct <= 5 and -16383 + 768i1 + i2 <= 16384ct <= 768i1 + i2 and 0 <= slot <= 16383 }">
#layout2 = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : (-768i0 - i1 + slot + 16384*floor((768i0 + i1)/16384)) mod 131072 = 0 and 0 <= i0 <= 127 and 0 <= i1 <= 767 and -16383 + 768i0 + i1 <= 16384ct <= 768i0 + i1 and 0 <= slot <= 16383 }">
module {
func.func @forward(%arg0: !secret.secret> {tensor_ext.layout = #layout1}) -> (!secret.secret> {tensor_ext.layout = #layout}) {
%0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #layout1}) {
^body(%input0: tensor<1x128x768xf32>):
%collapsed = tensor.collapse_shape %input0 [[0, 1], [2]] {tensor_ext.layout = #layout2} : tensor<1x128x768xf32> into tensor<128x768xf32>
%1 = tensor_ext.convert_layout %collapsed {from_layout = #layout2, tensor_ext.layout = #layout, to_layout = #layout} : tensor<128x768xf32>
secret.yield %1 : tensor<128x768xf32>
} -> (!secret.secret> {tensor_ext.layout = #layout})
return %0 : !secret.secret>
}
}
```

The layout conversion here is taking the row-major layout, which for a tensor<128x768> requires 6 ciphertexts, and converting that to a per-row ciphertext layout, requiring 128 layouts. When the `ConvertConverLayout` pattern runs, it lowers this to a remap operation and a shape changing extract slice

```
#layout = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : (-i1 - 256ct + slot) mod 1024 = 0 and 0 <= i1 <= 16383 and 0 <= ct <= 127 and -767 + 16384i0 + i1 <= 768ct <= 16384i0 + i1 and 0 <= slot <= 16383 }">
#layout1 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : i0 = 0 and (-768i1 - i2 + slot + 16384*floor((768i1 + i2)/16384)) mod 131072 = 0 and 0 <= i1 <= 127 and 0 <= i2 <= 767 and 0 <= ct <= 5 and -16383 + 768i1 + i2 <= 16384ct <= 768i1 + i2 and 0 <= slot <= 16383 }">
#layout2 = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = i0 and (-i1 + slot) mod 1024 = 0 and 0 <= i0 <= 127 and 0 <= i1 <= 767 and 0 <= slot <= 16383 }">
#original_type = #tensor_ext.original_type, layout = #layout1>
#original_type1 = #tensor_ext.original_type, layout = #layout2>
"builtin.module"() ({
"func.func"() <{arg_attrs = [{tensor_ext.original_type = #original_type}], function_type = (!secret.secret>) -> !secret.secret>, res_attrs = [{tensor_ext.original_type = #original_type1}], sym_name = "forward"}> ({
^bb0(%arg0: !secret.secret>):
%0 = "secret.generic"(%arg0) ({
^bb0(%arg1: tensor<6x16384xf32>):
%1 = "tensor_ext.remap"(%arg1) <{permutation = #layout}> : (tensor<6x16384xf32>) -> tensor<6x16384xf32>
%2 = "tensor.extract_slice"(%1) <{operandSegmentSizes = array, static_offsets = array, static_sizes = array, static_strides = array}> : (tensor<6x16384xf32>) -> tensor<128x16384xf32>
"secret.yield"(%2) : (tensor<128x16384xf32>) -> ()
}) : (!secret.secret>) -> !secret.secret>
"func.return"(%0) : (!secret.secret>) -> ()
}) : () -> ()
}) : () -> ()
```

BUT I think remap should have the output type be 128x16384... so remap needs to change types at some point, and remap right now expects the same input and output type. OR the input should be inserted into a larger 128 length tensor that is zero-padded, and then remapped.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.