Skip to content

Instantly share code, notes, and snippets.

@j2kun
Created January 21, 2025 21:26
Show Gist options
  • Save j2kun/75b676fe90d5ba1f2c67ad1a2409406b to your computer and use it in GitHub Desktop.
Save j2kun/75b676fe90d5ba1f2c67ad1a2409406b to your computer and use it in GitHub Desktop.
#map = affine_map<(d0, d1) -> (d0 * 16 + d1)>
#map1 = affine_map<(d0, d1) -> (d1)>
#map2 = affine_map<(d0) -> (d0)>
#map3 = affine_map<(d0, d1, d2) -> (d0, d2)>
#map4 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
module {
func.func @matmul(%arg0: !secret.secret<tensor<1x16xf32>> {tensor_ext.layout = #tensor_ext.layout<layout = (d0, d1) -> (d0 * 16 + d1)>}) -> (!secret.secret<tensor<1x16xf32>> {tensor_ext.layout = #tensor_ext.layout<layout = (d0, d1) -> (d1)>}) {
%cst = arith.constant dense<"0x5036CB3DE147C3BEE4A9393E47C021BE40F376BFA1078D3E8D53EB3DD6E0493EEFFC3CBFBEB947BE4597B5BBD185903E9B9C1BBEEB0713BD23B418BF66736C3EED693F3E584F72BF3CB9E83EBD0E8D3E4D87BDBE5A0439BFBE94AABECDCA91BE695FA93E870B93BE576920BF6294083F4C08633DCBACC6BDD8C9243FABF141BE647E8E3F27116D3DBA00FA3DDDD4C93EA96AA03E1FD4A7BE3C3297BD387D02BFA695923E3402CB3E6A4E0F3E8D0700BF195E3E3ECA2E0EBF6CAA17BE63FE853E1C7029BFAA05383F0DBFF0BEBF82CB3F8D9F843D3640A63F75BF7FBEF615053D1937CB3FC68B41BEEC66B9BED998223F90944FBD28F39CBC21AE853F26F7803FBEE1923F32109FBF79DF193F726237BFBF6FFB3F55B69FBF1EFEF83E7D4EB63F553A37BF50F054BF4072EE3FACE7A1BF511F9DBFE8A4C23F3C11793F68822ABF83F258BEC956703F4E00073FE1645E3F9C8203BF6D8B66BD1936893F0042113E6EC745BF161EB23E570AFE3D79CC633C8E44723F8D844E3FBE51FABE7BFA0F3F77DDA1BFADE2DCBE29DA10BEA4569FBD24B92ABE4DF072BFAE8AA13E4C661B3F3DCF823E4FDC1F3E961D7C3E1039C43D665C36BF2791AF3D47B452BE34128B3ED7EDB93CEE8F443E1C9693BDB0863BBF7F70C5BEC166A93EAF6CACBDBF5F023FEC98153FB562A1BD5FEE0ABD8FB21CBFCE54193FF31C79BE2A0A763EA3B655BF309239C03EF9903F12360BBFD15A1FC0733F6C3F8D4BDF3F615C2DC083868D3FC2D49C3FD9A115C01EC09EBF36CD1B3F9ABE19C05129963F4CA4BDBFAB2F1C3E842D833E90A07CBFEEFCFBBD97BE063E7CE7DA3EBF4AEABD473F593F497814BD8523E03DCCA8D6BDEA77253E99D43DBEECACB53E74A657BF7AE739BECC272F3FB0C3163F7803133F051EBDBE94A451BF1F83C13F9FD976BFE32D25BF911CAB3F07ED413DD65AFDBE532AA23F85E451BF88925B3E3D09BD3D6C0AC33F2B3A19BFD1DB2ABFCDB2E83EAE500E3FE4AA0B3F0284113EB809763EDF71D0BD4BC424BE13E9853ED757033F15A656BE522F40BEA19AA4BEF1F9953E0FDF2E3D198BD9BE5D971B3FDC29103F598190BF0C8726BEF339193FCD4C10BF382D6CBE7A020C3F016DA2BE590DF63E1923163D8B94383D1AD4EABEFDA50DBBBDA0BABCA75C0DBF1091853ECC190A3F369C0D3FA2AD41BE1B19E13CC88265BF2DD392BD6509A73ED73A6F3EC4280CBFC24284BE76727CBEC5DE023F79B7B8BE6E8E23BF12C739BDED900CBEECD986BF65AB74BF1A15F63EA1F5FA3E6B2BABBE9C4FECB895E499BE2D0268BE8EA7EABE374FFD3DADBB19BFC759C83F4A69D73E37B836BFE5F4E1BE19464BBFAE4E853D022E55BE1CDB073F6E31C9BEC202C5BE4BF853BEEE54DB3EEBC9613E74C317BEB9F2A3BE755B6F3F37CE383FF01E2F3D989532BE1C591EBE"> : tensor<16x16xf32>
%cst_0 = arith.constant dense<[-0.45141533, -0.0277900472, 0.311195374, 0.18254894, -0.258809537, 0.497506738, 0.00115649134, -0.194445714, 0.158549473, 0.000000e+00, 0.310650676, -0.214976981, -0.023661999, -0.392960966, 6.472870e-01, 0.831665277]> : tensor<16xf32>
%0 = secret.generic attrs = {layout = [#map]} {
%2 = tensor.empty() : tensor<16x16xf32>
secret.yield %2 : tensor<16x16xf32>
} -> !secret.secret<tensor<16x16xf32>>
%1 = secret.generic ins(%arg0, %0 : !secret.secret<tensor<1x16xf32>>, !secret.secret<tensor<16x16xf32>>) attrs = {arg0 = {layout = #map}, arg1 = {layout = #map}, layout = [#map1]} {
^body(%input0: tensor<1x16xf32>, %input1: tensor<16x16xf32>):
%collapsed = tensor.collapse_shape %input0 [[0, 1]] {layout = [#map2]} : tensor<1x16xf32> into tensor<16xf32>
%2 = linalg.vecmat {indexing_maps = [#map3, #map4, #map5], layout = [#map2]} ins(%collapsed, %cst : tensor<16xf32>, tensor<16x16xf32>) outs(%cst_0 : tensor<16xf32>) -> tensor<16xf32>
%expanded = tensor.expand_shape %2 [[0, 1]] output_shape [1, 16] {layout = [#map1]} : tensor<16xf32> into tensor<1x16xf32>
secret.yield %expanded : tensor<1x16xf32>
} -> !secret.secret<tensor<1x16xf32>>
return %1 : !secret.secret<tensor<1x16xf32>>
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment