wip: rotom seeding layout pass

[copybara-service]

wip: rotom seeding layout pass

This pass assigns an initial layout to SSA values in a module for a Rotom optimizer to ingest and optimize. The initial layouts traverse over all combinations of 1 <= i <= h and 1 <= j <= w, where i * j == n, and i and j are also powers-of-two.

For example, for the 4x4 tensors below: it creates the following 4 layouts corresponding to

[0:2:2];[0:2:1][1:4:1]
[0:2:2];[1:4:1][0:2:1]
[1:2:2];[1:2:1][0:4:1]
[1:2:2];[0:4:1][1:2:1]

#layout = #rotom.layout<dims = [#rotom.dim<dim = 0, size = 2, stride = 2>, #rotom.dim<dim = 0, size = 2>, #rotom.dim<dim = 1, size = 4>], n = 8>
#layout1 = #rotom.layout<dims = [#rotom.dim<dim = 0, size = 2, stride = 2>, #rotom.dim<dim = 1, size = 4>, #rotom.dim<dim = 0, size = 2>], n = 8>
#layout2 = #rotom.layout<dims = [#rotom.dim<dim = 1, size = 2, stride = 2>, #rotom.dim<dim = 0, size = 4>, #rotom.dim<dim = 1, size = 2>], n = 8>
#layout3 = #rotom.layout<dims = [#rotom.dim<dim = 1, size = 2, stride = 2>, #rotom.dim<dim = 1, size = 2>, #rotom.dim<dim = 0, size = 4>], n = 8>
#seed = #rotom.seed<layouts = [#layout, #layout1, #layout2, #layout3]>
module {
  func.func @test_seeding(%arg0: !secret.secret<tensor<4x4xf32>> {rotom.seed = #seed}, %arg1: tensor<4x4xf32> {rotom.seed = #seed}) -> !secret.secret<tensor<4x4xf32>> {
    %0 = secret.generic(%arg0: !secret.secret<tensor<4x4xf32>> {rotom.seed = #seed}) {
    ^body(%input0: tensor<4x4xf32>):
      %1 = arith.addf %input0, %arg1 {rotom.seed = #seed} : tensor<4x4xf32>
      secret.yield %1 : tensor<4x4xf32>
    } -> !secret.secret<tensor<4x4xf32>>
    return %0 : !secret.secret<tensor<4x4xf32>>
  }
}

[asraa]

@edwjchen i started the skeleton of the pass here! does it roughly look right? Two things:

• I still have to add in padding since not all tensors will be power of two sized
• I'll also add in caching so I don't have to recompute seeds for tensors of the same size
• Do I need to also seed the results? I noticed in the Rotom repository that gen_tensor will initialize all the created tensors, but there's also functions to generate the layouts for results of adds/muls/etc.