IntegerRelation computation: confusing composition of relations?

[asraa]

Suppose I have the following IR:

#layout = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-4i0 - i1 + slot) mod 16 = 0 and 0 <= i0 <= 3 and 0 <= i1 <= 3 and 0 <= slot <= 1023 }">
#layout1 = #tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : i0 = 0 and i1 = 0 and ct = 0 and (-4i2 - i3 + slot) mod 16 = 0 and 0 <= i2 <= 3 and 0 <= i3 <= 3 and 0 <= slot <= 1023 }">
module {
  func.func @main(%arg0: !secret.secret<tensor<1x1x4x4xf32>> {tensor_ext.layout = #layout1}, %arg1: tensor<2x1x3x3xf32>) -> (!secret.secret<tensor<4x4xf32>> {tensor_ext.layout = #layout}) {
    %0 = secret.generic(%arg0: !secret.secret<tensor<1x1x4x4xf32>> {tensor_ext.layout = #layout1}) {
    ^body(%input0: tensor<1x1x4x4xf32>):
      %extracted_slice = tensor.extract_slice %input0[0, 0, 0, 0] [1, 1, 4, 4] [1, 1, 1, 1] {tensor_ext.layout = #layout} : tensor<1x1x4x4xf32> to tensor<4x4xf32>
      secret.yield %extracted_slice : tensor<4x4xf32>
    } -> (!secret.secret<tensor<4x4xf32>> {tensor_ext.layout = #layout})
    return %0 : !secret.secret<tensor<4x4xf32>>
  }
}

The input layout #layout1 maps from tensor<1x1x4x4> -> (ct, slot) by ordering the 4x4 into a row major layout (it repeats the 16 elements in the 1024 sized ciphertext). The result layout of the slice is the same with the first two dimensions eliminated, and maps tensor<4x4> -> (ct, slot).

If we want to compute the kernel for a tensor.extract_slice, then in this particular case we would have a no-op - the input and output layouts are identical within the tensor<1x1024> tensor-semantic shape. But in generality, we could compute the kernel as a tensor_ext.remap operation that converts between the (ct, slot) representation of the tensor<1x1x4x4> to the slice extracted in the ciphertext domain. This can be done by (1) invert #layout to go from (ct, slot) -> tensor<1x1x4x4> (2) compose with the relation that extracts a slice tensor<1x1x4x4> -> tensor<4x4> (3) compose with the ciphertext semantic relation of the slice from tensor<4x4> -> (ct, slot).

For these four steps, we have:

#invert = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot, o2, o3] : i0 = 0 and ct = 0 and slot = 0 and (-i1 + 4o2 + o3) mod 16 = 0 and 0 <= i1 <= 1023 and 0 <= o2 <= 3 and 0 <= o3 <= 3 }">
#invert_extract = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and (-i1 + 4ct + slot) mod 16 = 0 and 0 <= i1 <= 1023 and 0 <= ct <= 3 and 0 <= slot <= 3 }">
#invert_extract_map = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : exists (e0, e1, e2: i0 = 0 and ct = 0 and 16e2 = -i1 + slot + 16e0 and 0 <= i1 <= 1023 and 0 <= slot <= 1023 and 0 <= e1 <= 3 and -3 + i1 - 16e0 <= 4e1 <= i1 - 16e0) }">

We expect that the final result #invert_extract_map be the identity, but now it has quantifiers (including when we simplify). I think there is something interesting that happens because of the slot repetition...

[asraa]

OK the codegen for this relation is quick and simple and looks correct (minus why it's not just S(0, c1, 0, c1))

or (int c1 = 0; c1 <= 1023; c1 += 1)\n  for (int c3 = c1 % 16; c3 <= 1023; c3 += 16)\n    S(0, c3, 0, c1);

But the shift network does not :|