## float32[64,128]：数值检查通过
cfence: 1
dma.simple: 2
sadd.s32: 2
sand.u32: 3
sbr.rel: 1
simm.s32: 1
sld: 3
sne.s32: 1
sor.u32: 2
sshll.u32: 2
sshrl.u32: 1
vld.8x128: 8
vmul.8x128.f32: 8
vst.8x128: 8
vsyncadd.remote.s32: 1
vsyncadd.s32: 3
vtrace: 2
vwait.ge: 3
# function jit(scale)/scale/mpmd_map [symbol 4, scale.1]; image bundles [503, 507)
{ misc: vtrace 0x80000000 }
# entry bundle: %scale.1 = custom-call(%x_hbm.1)
{ s1: sld s6, [smem:0x1] }
{ s0: sne.s32 p0, s6, 0 }
{ s0: @p0 sbr.rel L_020d }
# end function jit(scale)/scale/mpmd_map [symbol 4, scale.1]
{}
# function jit(scale)/scale/mpmd_map [symbol 4, scale.1]; image bundles [508, 525)
{ s0: simm.s32 s7, 0 ;  # materialized constant
  .encoding { s0.y = lo(imm0) } }
{ s0: dma.simple [vmem:s7], [hbm:s0], length=64, dst_flag=[sflag:52] }
{ misc: vwait.ge [sflag:52], 64 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:38:16-38:59 [dma_wait; LLO 62]
  .encoding { misc.value = imm2 } }
{ vld: vld.8x128 v0, [vmem:0x0] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 25]
  misc: vsyncadd.s32 [sflag:52], -64 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:38:16-38:59 [dma_wait; LLO 63]
  .encoding { misc.value = imm2 } }
{ va0: vmul.8x128.f32 v1, 2.0, v0 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 33]
  vld: vld.8x128 v2, [vmem:0x8] }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 26]
{ va0: vmul.8x128.f32 v3, 2.0, v2 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 34]
  vld: vld.8x128 v4, [vmem:0x10] }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 27]
{ va0: vmul.8x128.f32 v5, 2.0, v4 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 35]
  vst: vst.8x128 [vmem:0x0], v1 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 41]
  vld: vld.8x128 v6, [vmem:0x18] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 28]
  .encoding { vst.offset = imm2 } }
{ va0: vmul.8x128.f32 v7, 2.0, v6 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 36]
  vst: vst.8x128 [vmem:0x8], v3 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 42]
  vld: vld.8x128 v8, [vmem:0x20] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 29]
  .encoding { vst.offset = imm2 } }
{ va0: vmul.8x128.f32 v9, 2.0, v8 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 37]
  vst: vst.8x128 [vmem:0x10], v5 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 43]
  vld: vld.8x128 v10, [vmem:0x28] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 30]
  .encoding { vst.offset = imm2 } }
{ va0: vmul.8x128.f32 v11, 2.0, v10 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 38]
  vst: vst.8x128 [vmem:0x18], v7 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 44]
  vld: vld.8x128 v12, [vmem:0x30] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 31]
  .encoding { vst.offset = imm2 } }
{ va0: vmul.8x128.f32 v13, 2.0, v12 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 39]
  vst: vst.8x128 [vmem:0x20], v9 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 45]
  vld: vld.8x128 v14, [vmem:0x38] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 32]
  .encoding { vst.offset = imm2 } }
{ va0: vmul.8x128.f32 v15, 2.0, v14 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 40]
  vst: vst.8x128 [vmem:0x28], v11 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 46]
{ vst: vst.8x128 [vmem:0x30], v13 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 47]
{ vst: vst.8x128 [vmem:0x38], v15 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 48]
{ s0: dma.simple [hbm:s1], [vmem:s7], length=64, dst_flag=[sflag:52] }
{ misc: vwait.ge [sflag:52], 64 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:43:16-43:59 [dma_wait; LLO 64]
  .encoding { misc.value = imm2 } }
{ misc: vsyncadd.s32 [sflag:52], -64 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:43:16-43:59 [dma_wait; LLO 65]
  .encoding { misc.value = imm2 } }
# end function jit(scale)/scale/mpmd_map [symbol 4, scale.1]
# exit bundle: %scale.1 = custom-call(%x_hbm.1)
# ... 省略 runtime 代码 1 个 bundle
# function jit(scale)/scale/mpmd_map [symbol 4, scale.1]; image bundles [526, 542)
{ s1: sld s20, [smem:0x0] }
{ s1: sld s21, [smem:0x1] }
{ s0: sand.u32 s24, 0xfff, s20 }
{ s0: sadd.s32 s22, 1, s21 ;
  s1: sshll.u32 s27, s24, 0x12 }
{ s0: sand.u32 s23, 0x1, s22 }  # smod.u32 w/div 2
{ s0: sand.u32 s25, 0x3, s23 }
{ s0: sshll.u32 s26, s25, 0x10 }
{ s0: sadd.s32 s28, 131072, s26 }
{ s0: sshrl.u32 s29, s28, 0x2 }
{ s0: sor.u32 s30, 0x2d, s29 }
{ s0: sor.u32 s2, s30, s27 }
{ misc: vsyncadd.remote.s32 [sflag:s2], 1 }
{ misc: vwait.ge [sflag:45], 1 ;
  .encoding { misc.value = imm2 } }
{ misc: vsyncadd.s32 [sflag:45], -1 ;
  .encoding { misc.value = imm2 } }
{ misc: cfence }
{ misc: vtrace 0x90000000 }
# end function jit(scale)/scale/mpmd_map [symbol 4, scale.1]

## bfloat16[64,128]：数值检查通过
cfence: 1
dma.simple: 2
sadd.s32: 2
sand.u32: 3
sbr.rel: 1
simm.s32: 1
sld: 3
sne.s32: 1
sor.u32: 2
sshll.u32: 2
sshrl.u32: 1
vld.8x128: 4
vmul.8x128.f32: 8
vpackc.8x128.f32.f16: 4
vst.8x128: 4
vsyncadd.remote.s32: 1
vsyncadd.s32: 3
vtrace: 2
vunpackl.8x128.f16.f32: 4
vunpacku.8x128.f16.f32: 4
vwait.ge: 3
# function jit(scale)/scale/mpmd_map [symbol 4, scale.1]; image bundles [503, 507)
{ misc: vtrace 0x80000000 }
# entry bundle: %scale.1 = custom-call(%x_hbm.1)
{ s1: sld s6, [smem:0x1] }
{ s0: sne.s32 p0, s6, 0 }
{ s0: @p0 sbr.rel L_020f }
# end function jit(scale)/scale/mpmd_map [symbol 4, scale.1]
{}
# function jit(scale)/scale/mpmd_map [symbol 4, scale.1]; image bundles [508, 527)
{ s0: simm.s32 s7, 0 ;  # materialized constant
  .encoding { s0.y = lo(imm0) } }
{ s0: dma.simple [vmem:s7], [hbm:s0], length=32, dst_flag=[sflag:52] }
{ misc: vwait.ge [sflag:52], 32 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:38:16-38:59 [dma_wait; LLO 120]
  .encoding { misc.value = imm2 } }
{ vld: vld.8x128 v0, [vmem:0x0] ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 78] | Coalesced load.
  misc: vsyncadd.s32 [sflag:52], -32 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:38:16-38:59 [dma_wait; LLO 121]
  .encoding { misc.value = imm2 } }
{ va0: vunpackl.8x128.f16.f32 v1, v0 ;
  va1: vunpacku.8x128.f16.f32 v2, v0 ;
  vld: vld.8x128 v3, [vmem:0x8] }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 113] | Coalesced load.
{ va0: vunpackl.8x128.f16.f32 v4, v3 ;
  va1: vunpacku.8x128.f16.f32 v5, v3 ;
  vld: vld.8x128 v6, [vmem:0x10] }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 114] | Coalesced load.
{ va0: vmul.8x128.f32 v7, 2.0, v1 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 41]
  va1: vunpackl.8x128.f16.f32 v8, v6 ;
  vld: vld.8x128 v9, [vmem:0x18] }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:45 [get; LLO 115] | Coalesced load.
{ va0: vmul.8x128.f32 v10, 2.0, v2 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 42]
  va1: vunpacku.8x128.f16.f32 v11, v6 }
{ va0: vmul.8x128.f32 v12, 2.0, v4 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 43]
  va1: vunpackl.8x128.f16.f32 v13, v9 }
{ va0: vmul.8x128.f32 v15, 2.0, v5 ;
  va1: vpackc.8x128.f32.f16 v14, v10, v7 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 44]
{ va0: vmul.8x128.f32 v16, 2.0, v8 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 45]
  va1: vunpacku.8x128.f16.f32 v17, v9 }
{ va0: vmul.8x128.f32 v19, 2.0, v11 ;
  va1: vpackc.8x128.f32.f16 v18, v15, v12 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 46]
  vst: vst.8x128 [vmem:0x0], v14 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 97] | Coalesced store.
  .encoding { vst.opcode = s0 } }
{ va0: vmul.8x128.f32 v20, 2.0, v13 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 47]
{ va0: vmul.8x128.f32 v22, 2.0, v17 ;
  va1: vpackc.8x128.f32.f16 v21, v19, v16 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:34-42:49 [mul; LLO 48]
  vst: vst.8x128 [vmem:0x8], v18 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 116] | Coalesced store.
{ va0: vpackc.8x128.f32.f16 v23, v22, v20 ;
  vst: vst.8x128 [vmem:0x10], v21 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 117] | Coalesced store.
{ vst: vst.8x128 [vmem:0x18], v23 }  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:42:20-42:31 [swap; LLO 118] | Coalesced store.
{ s0: dma.simple [hbm:s1], [vmem:s7], length=32, dst_flag=[sflag:52] }
{ misc: vwait.ge [sflag:52], 32 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:43:16-43:59 [dma_wait; LLO 122]
  .encoding { misc.value = imm2 } }
{ misc: vsyncadd.s32 [sflag:52], -32 ;  # chapter01/05_vector_layout/01_pallas_scale_64x128.py:43:16-43:59 [dma_wait; LLO 123]
  .encoding { misc.value = imm2 } }
# end function jit(scale)/scale/mpmd_map [symbol 4, scale.1]
# exit bundle: %scale.1 = custom-call(%x_hbm.1)
# ... 省略 runtime 代码 1 个 bundle
# function jit(scale)/scale/mpmd_map [symbol 4, scale.1]; image bundles [528, 544)
{ s1: sld s20, [smem:0x0] }
{ s1: sld s21, [smem:0x1] }
{ s0: sand.u32 s24, 0xfff, s20 }
{ s0: sadd.s32 s22, 1, s21 ;
  s1: sshll.u32 s27, s24, 0x12 }
{ s0: sand.u32 s23, 0x1, s22 }  # smod.u32 w/div 2
{ s0: sand.u32 s25, 0x3, s23 }
{ s0: sshll.u32 s26, s25, 0x10 }
{ s0: sadd.s32 s28, 131072, s26 }
{ s0: sshrl.u32 s29, s28, 0x2 }
{ s0: sor.u32 s30, 0x2d, s29 }
{ s0: sor.u32 s2, s30, s27 }
{ misc: vsyncadd.remote.s32 [sflag:s2], 1 }
{ misc: vwait.ge [sflag:45], 1 ;
  .encoding { misc.value = imm2 } }
{ misc: vsyncadd.s32 [sflag:45], -1 ;
  .encoding { misc.value = imm2 } }
{ misc: cfence }
{ misc: vtrace 0x90000000 }
# end function jit(scale)/scale/mpmd_map [symbol 4, scale.1]

## int8[64,128]：编译失败
INTERNAL: Mosaic failed to compile TPU kernel with MLIR module name `scale` HLO name `scale.1`: Not implemented: Only vector<i16> and vector<i32> are supported, but got 'i8'. Please cast your input.