module_forward_dispatch_9.mlir

hal.executable public @forward_dispatch_9 {
  hal.executable.variant public @embedded_elf_x86_64 target(<"llvm-cpu", "embedded-elf-x86_64", {cpu = "haswell", cpu_features = "-prfchw,-cldemote,+avx,+aes,+sahf,+pclmul,-xop,+crc32,-xsaves,-avx512fp16,-usermsr,-sm4,+sse4.1,-avx512ifma,+xsave,-avx512pf,+sse4.2,-tsxldtrk,-ptwrite,-widekl,-sm3,+invpcid,+64bit,-xsavec,-avx10.1-512,-avx512vpopcntdq,+cmov,-avx512vp2intersect,-avx512cd,+movbe,-avxvnniint8,-avx512er,-amx-int8,-kl,-avx10.1-256,-sha512,-avxvnni,-rtm,-adx,+avx2,-hreset,-movdiri,-serialize,-vpclmulqdq,-avx512vl,-uintr,-clflushopt,-raoint,-cmpccxadd,+bmi,-amx-tile,+sse,-gfni,-avxvnniint16,-amx-fp16,+xsaveopt,+rdrnd,-avx512f,-amx-bf16,-avx512bf16,-avx512vnni,+cx8,-avx512bw,+sse3,-pku,+fsgsbase,-clzero,-mwaitx,-lwp,+lzcnt,-sha,-movdir64b,-wbnoinvd,-enqcmd,-prefetchwt1,-avxneconvert,-tbm,-pconfig,-amx-complex,+ssse3,+cx16,+bmi2,+fma,+popcnt,-avxifma,+f16c,-avx512bitalg,-rdpru,-clwb,+mmx,+sse2,-rdseed,-avx512vbmi2,-prefetchi,-rdpid,-fma4,-avx512vbmi,-shstk,-vaes,-waitpkg,-sgx,+fxsr,-avx512dq,-sse4a", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", native_vector_size = 32 : index, target_triple = "x86_64-unknown-unknown-eabi-elf", ukernels = "--iree-llvmcpu-stack-allocation-limit=256000"}>) {
    hal.executable.export public @forward_dispatch_9_generic_4x32x32_f16 ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer, ReadOnly>, <1, storage_buffer, ReadOnly>, <2, storage_buffer>]>]>) {
    ^bb0(%arg0: !hal.device loc("aten::sub"("<eval_with_key>.5":67:10))):
      %x, %y, %z = flow.dispatch.workgroup_count_from_slice  loc("aten::sub"("<eval_with_key>.5":67:10))
      hal.return %x, %y, %z : index, index, index loc("aten::sub"("<eval_with_key>.5":67:10))
    } loc("aten::sub"("<eval_with_key>.5":67:10))
    builtin.module {
      func.func @forward_dispatch_9_generic_4x32x32_f16() {
        %c106496 = arith.constant 106496 : index loc("aten::view"("<eval_with_key>.5":51:13))
        %c3895398400 = arith.constant 3895398400 : index loc(fused["aten::select"("<eval_with_key>.5":54:15), "aten::select"("<eval_with_key>.5":56:15)])
        %c3895398656 = arith.constant 3895398656 : index loc(fused["aten::select"("<eval_with_key>.5":54:15), "aten::select"("<eval_with_key>.5":56:15)])
        %c65536 = arith.constant 65536 : index loc("aten::sub"("<eval_with_key>.5":67:10))
        %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c106496) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<4x1x32x32x2xf16>> loc("aten::view"("<eval_with_key>.5":51:13))
        %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c3895398400) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<4x32xf16>> loc("aten::select"("<eval_with_key>.5":54:15))
        %2 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c3895398656) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<4x32xf16>> loc("aten::select"("<eval_with_key>.5":56:15))
        %3 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c65536) : !flow.dispatch.tensor<writeonly:tensor<4x32x32xf16>> loc("aten::sub"("<eval_with_key>.5":67:10))
        %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [4, 32], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<4x32xf16>> -> tensor<4x32xf16> loc("aten::sub"("<eval_with_key>.5":67:10))
        %5 = flow.dispatch.tensor.load %2, offsets = [0, 0], sizes = [4, 32], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<4x32xf16>> -> tensor<4x32xf16> loc("aten::sub"("<eval_with_key>.5":67:10))
        %6 = tensor.empty() : tensor<4x32x32xf16> loc("aten::mul"("<eval_with_key>.5":65:12))
        %7 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0, 1], sizes = [4, 1, 32, 32, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<4x1x32x32x2xf16>> -> tensor<4x32x32xf16> loc("aten::select"("<eval_with_key>.5":55:15))
        %8 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0, 0], sizes = [4, 1, 32, 32, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<4x1x32x32x2xf16>> -> tensor<4x32x32xf16> loc("aten::select"("<eval_with_key>.5":53:13))
        %9 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%8, %4, %7, %5 : tensor<4x32x32xf16>, tensor<4x32xf16>, tensor<4x32x32xf16>, tensor<4x32xf16>) outs(%6 : tensor<4x32x32xf16>) {
        ^bb0(%in: f16 loc("aten::select"("<eval_with_key>.5":53:13)), %in_0: f16 loc("aten::select"("<eval_with_key>.5":54:15)), %in_1: f16 loc("aten::select"("<eval_with_key>.5":55:15)), %in_2: f16 loc("aten::select"("<eval_with_key>.5":56:15)), %out: f16 loc("aten::sub"("<eval_with_key>.5":67:10))):
          %10 = arith.mulf %in_1, %in_2 : f16 loc("aten::mul"("<eval_with_key>.5":66:12))
          %11 = arith.mulf %in, %in_0 : f16 loc("aten::mul"("<eval_with_key>.5":65:12))
          %12 = arith.subf %11, %10 : f16 loc("aten::sub"("<eval_with_key>.5":67:10))
          linalg.yield %12 : f16 loc("aten::sub"("<eval_with_key>.5":67:10))
        } -> tensor<4x32x32xf16> loc("aten::sub"("<eval_with_key>.5":67:10))
        flow.dispatch.tensor.store %9, %3, offsets = [0, 0, 0], sizes = [4, 32, 32], strides = [1, 1, 1] : tensor<4x32x32xf16> -> !flow.dispatch.tensor<writeonly:tensor<4x32x32xf16>> loc("aten::sub"("<eval_with_key>.5":67:10))
        return loc("aten::sub"("<eval_with_key>.5":67:10))
      } loc("aten::sub"("<eval_with_key>.5":67:10))
    } loc("aten::sub"("<eval_with_key>.5":67:10))
  } loc("aten::sub"("<eval_with_key>.5":67:10))
} loc("aten::sub"("<eval_with_key>.5":67:10))