im2col.mlir

func.func @img2col(%arg0: tensor<128x1026x1026xf32>) -> tensor<128x3x3x1024x1024xbf16> {
  %0 = tensor.empty() : tensor<128x3x3x1024x1024xbf16>
  %c1 = arith.constant 1 : index
  %c0 = arith.constant 0 : index
  %cst = arith.constant 0.000000e+00 : f32
  %c128 = arith.constant 128 : index
  %c1024 = arith.constant 1024 : index
  %1 = scf.for %arg1 = %c0 to %c128 step %c1 iter_args(%arg2 = %0) -> (tensor<128x3x3x1024x1024xbf16>) {
    %2 = scf.for %arg3 = %c0 to %c1024 step %c1 iter_args(%arg4 = %arg2) -> (tensor<128x3x3x1024x1024xbf16>) {
      %3 = scf.for %arg5 = %c0 to %c1024 step %c1 iter_args(%arg6 = %arg4) -> (tensor<128x3x3x1024x1024xbf16>) {
        %4 = tensor.extract_slice %arg0[%arg1, %arg3, %arg5][1, 3, 3][1, 1, 1] : tensor<128x1026x1026xf32> to tensor<1x3x3xf32>
        %5 = arith.truncf %4 : tensor<1x3x3xf32> to tensor<1x3x3xbf16>
        %6 = tensor.insert_slice %5 into %arg6[%arg1, %c0, %c0, %arg3, %arg5][1, 3, 3, 1, 1][1, 1, 1, 1, 1] : tensor<1x3x3xbf16> into tensor<128x3x3x1024x1024xbf16>
        scf.yield %6 : tensor<128x3x3x1024x1024xbf16>
      }
      scf.yield %3 : tensor<128x3x3x1024x1024xbf16>
    }
    scf.yield %2 : tensor<128x3x3x1024x1024xbf16>
  }
  return %1 : tensor<128x3x3x1024x1024xbf16>
}