attnbf16.mlir

#map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d3)>
#map3 = affine_map<(d0, d1, d2, d3, d4, d5) -> ()>
#map4 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
module {
  func.func @attention(%arg0: tensor<1x128x32x64xbf16>, %arg1: tensor<1x128x32x64xbf16>, %arg2: tensor<1x128x32x64xbf16>, %arg3: tensor<1x128x32x64xbf16>) -> tensor<1x128x32x64xbf16> {
    %cst = arith.constant 1.250000e-01 : bf16
    %0 = tensor.empty() : tensor<1x128x32x64xbf16>
    %1 = iree_linalg_ext.attention {indexing_maps = [#map, #map1, #map2, #map3, #map4]} ins(%arg0, %arg1, %arg2, %cst : tensor<1x128x32x64xbf16>, tensor<1x128x32x64xbf16>, tensor<1x128x32x64xbf16>, bf16) outs(%0 : tensor<1x128x32x64xbf16>) {
    ^bb0(%arg4: f32):
      iree_linalg_ext.yield %arg4 : f32
    } -> tensor<1x128x32x64xbf16>
    return %1 : tensor<1x128x32x64xbf16>
  }
}