tensor_ext.remap doesn't handle shape expansion properly
- 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
Assessment
This issue has not been assessed yet.