// Runs with the integer rewrite enabled and disabled.
// Inputs stay unfolded so these checks execute the compiled contractions.

// Both correction products and the cross term vary by block and output axis.
func.func @blockwise_zero_point_corrections() {
  %aq = util.unfoldable_constant dense<[[[2, 3], [4, 6]], [[5, 8], [9, 10]]]> : tensor<2x2x2xi8>
  %bq = util.unfoldable_constant dense<[[[2, 4, 6], [3, 5, 8]], [[6, 7, 9], [7, 9, 10]]]> : tensor<2x2x3xi8>
  %sa = util.unfoldable_constant dense<[[1.0, 0.5], [2.0, 1.0]]> : tensor<2x2xf32>
  %sb = util.unfoldable_constant dense<[[1.0, 2.0, 1.0], [0.5, 1.0, 2.0]]> : tensor<2x3xf32>
  %za = util.unfoldable_constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi8>
  %zb = util.unfoldable_constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi8>
  %ai = tensor.empty() : tensor<2x2x2xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, g, k) -> (m, g, k)>,
                        affine_map<(m, g, k) -> (m, g)>,
                        affine_map<(m, g, k) -> (m, g)>,
                        affine_map<(m, g, k) -> (m, g, k)>]>
      ins(%aq, %sa, %za : tensor<2x2x2xi8>, tensor<2x2xf32>, tensor<2x2xi8>)
      outs(%ai : tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
  %bi = tensor.empty() : tensor<2x2x3xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(g, k, n) -> (g, k, n)>,
                        affine_map<(g, k, n) -> (g, n)>,
                        affine_map<(g, k, n) -> (g, n)>,
                        affine_map<(g, k, n) -> (g, k, n)>]>
      ins(%bq, %sb, %zb : tensor<2x2x3xi8>, tensor<2x3xf32>, tensor<2x3xi8>)
      outs(%bi : tensor<2x2x3xf32>) -> tensor<2x2x3xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<2x3xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<2x3xf32>) -> tensor<2x3xf32>
  %result = linalg.contract
      indexing_maps = [affine_map<(m, n, g, k) -> (m, g, k)>,
                       affine_map<(m, n, g, k) -> (g, k, n)>,
                       affine_map<(m, n, g, k) -> (m, n)>]
      ins(%a, %b : tensor<2x2x2xf32>, tensor<2x2x3xf32>)
      outs(%init : tensor<2x3xf32>) -> tensor<2x3xf32>
  check.expect_eq_const(%result, dense<[[9.0, 26.0, 35.0], [38.0, 110.0, 140.0]]> : tensor<2x3xf32>) : tensor<2x3xf32>
  return
}

// Nonuniform image values distinguish the strided/dilated windows. Each output
// channel has a different zero point and scale but the same dequantized filter.
func.func @convolution_zero_point_corrections() {
  %aq = util.unfoldable_constant dense<[[[[1], [2], [3], [4], [5]], [[6], [7], [8], [9], [10]], [[11], [12], [13], [14], [15]], [[16], [17], [18], [19], [20]], [[21], [22], [23], [24], [25]]]]> : tensor<1x5x5x1xi8>
  %bq = util.unfoldable_constant dense<[[[[7, 9]], [[7, 9]]], [[[7, 9]], [[7, 9]]]]> : tensor<2x2x1x2xi8>
  %sa = util.unfoldable_constant 0.5 : f32
  %sb = util.unfoldable_constant dense<[1.0, 0.5]> : tensor<2xf32>
  %za = util.unfoldable_constant 3 : i8
  %zb = util.unfoldable_constant dense<[6, 7]> : tensor<2xi8>
  %ai = tensor.empty() : tensor<1x5x5x1xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(n, h, w, c) -> (n, h, w, c)>,
                        affine_map<(n, h, w, c) -> ()>,
                        affine_map<(n, h, w, c) -> ()>,
                        affine_map<(n, h, w, c) -> (n, h, w, c)>]>
      ins(%aq, %sa, %za : tensor<1x5x5x1xi8>, f32, i8)
      outs(%ai : tensor<1x5x5x1xf32>) -> tensor<1x5x5x1xf32>
  %bi = tensor.empty() : tensor<2x2x1x2xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(h, w, c, f) -> (h, w, c, f)>,
                        affine_map<(h, w, c, f) -> (f)>,
                        affine_map<(h, w, c, f) -> (f)>,
                        affine_map<(h, w, c, f) -> (h, w, c, f)>]>
      ins(%bq, %sb, %zb : tensor<2x2x1x2xi8>, tensor<2xf32>, tensor<2xi8>)
      outs(%bi : tensor<2x2x1x2xf32>) -> tensor<2x2x1x2xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x2x2x2xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  %result = linalg.conv_2d_nhwc_hwcf {dilations = dense<2> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
      ins(%a, %b : tensor<1x5x5x1xf32>, tensor<2x2x1x2xf32>)
      outs(%init : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  check.expect_eq_const(%result, dense<[[[[8.0, 8.0], [12.0, 12.0]], [[28.0, 28.0], [32.0, 32.0]]]]> : tensor<1x2x2x2xf32>) : tensor<1x2x2x2xf32>
  return
}

// NCHW layout with two input channels: all three filter dimensions reduce.
// Distinct signed data and per-filter parameters expose channel/layout swaps.
func.func @convolution_nchw_asymmetric() {
  %aq = util.unfoldable_constant dense<[[[[-8, -5, -2, 1], [-3, 0, 3, 6], [2, 5, 8, -6]], [[3, 6, -8, -5], [8, -6, -3, 0], [-4, -1, 2, 5]]]]> : tensor<1x2x3x4xi8>
  %bq = util.unfoldable_constant dense<[[[[-5, -4], [-2, -1]], [[0, 1], [3, 4]]], [[[2, 3], [5, -5]], [[-4, -3], [-1, 0]]]]> : tensor<2x2x2x2xi8>
  %sb = util.unfoldable_constant dense<[1.0, 0.25]> : tensor<2xf32>
  %zb = util.unfoldable_constant dense<[-1, 2]> : tensor<2xi8>
  %sa = util.unfoldable_constant 0.5 : f32
  %za = util.unfoldable_constant -2 : i8
  %ai = tensor.empty() : tensor<1x2x3x4xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(n, c, h, w) -> (n, c, h, w)>,
                        affine_map<(n, c, h, w) -> ()>,
                        affine_map<(n, c, h, w) -> ()>,
                        affine_map<(n, c, h, w) -> (n, c, h, w)>]>
      ins(%aq, %sa, %za : tensor<1x2x3x4xi8>, f32, i8)
      outs(%ai : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32>
  %bi = tensor.empty() : tensor<2x2x2x2xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(f, c, h, w) -> (f, c, h, w)>,
                        affine_map<(f, c, h, w) -> (f)>,
                        affine_map<(f, c, h, w) -> (f)>,
                        affine_map<(f, c, h, w) -> (f, c, h, w)>]>
      ins(%bq, %sb, %zb : tensor<2x2x2x2xi8>, tensor<2xf32>, tensor<2xi8>)
      outs(%bi : tensor<2x2x2x2xf32>) -> tensor<2x2x2x2xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x2x2x3xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
  %result = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
      ins(%a, %b : tensor<1x2x3x4xf32>, tensor<2x2x2x2xf32>)
      outs(%init : tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
  check.expect_eq_const(%result, dense<[[[[37.5, -7.5, -10.0], [-3.5, -6.0, 0.0]], [[-14.0, -4.125, 1.5], [-8.875, -3.25, 4.5]]]]> : tensor<1x2x2x3xf32>) : tensor<1x2x2x3xf32>
  return
}

// Channels remain parallel in depthwise convolution. Both operands have
// channel-specific scales and zero points; activation bytes exercise extui.
// Different stride/dilation axes distinguish the sampled spatial windows.
func.func @depthwise_asymmetric_strided_dilated() {
  %aq = util.unfoldable_constant dense<[[[[128, 133], [131, 136], [134, 139], [137, 142], [140, 145]], [[135, 140], [138, 143], [141, 146], [144, 130], [128, 133]], [[142, 128], [145, 131], [129, 134], [132, 137], [135, 140]], [[130, 135], [133, 138], [136, 141], [139, 144], [142, 128]]]]> : tensor<1x4x5x2xi8>
  %bq = util.unfoldable_constant dense<[[[-3, 4], [2, -1]], [[5, 0], [-2, 6]]]> : tensor<2x2x2xi8>
  %sa = util.unfoldable_constant dense<[0.5, 2.0]> : tensor<2xf32>
  %za = util.unfoldable_constant dense<[132, 136]> : tensor<2xi8>
  %sb = util.unfoldable_constant dense<[2.0, 0.25]> : tensor<2xf32>
  %zb = util.unfoldable_constant dense<[-1, 2]> : tensor<2xi8>
  %ai = tensor.empty() : tensor<1x4x5x2xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, zp_unsigned, indexing_maps = [affine_map<(n, h, w, c) -> (n, h, w, c)>,
                        affine_map<(n, h, w, c) -> (c)>,
                        affine_map<(n, h, w, c) -> (c)>,
                        affine_map<(n, h, w, c) -> (n, h, w, c)>]>
      ins(%aq, %sa, %za : tensor<1x4x5x2xi8>, tensor<2xf32>, tensor<2xi8>)
      outs(%ai : tensor<1x4x5x2xf32>) -> tensor<1x4x5x2xf32>
  %bi = tensor.empty() : tensor<2x2x2xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(h, w, c) -> (h, w, c)>,
                        affine_map<(h, w, c) -> (c)>,
                        affine_map<(h, w, c) -> (c)>,
                        affine_map<(h, w, c) -> (h, w, c)>]>
      ins(%bq, %sb, %zb : tensor<2x2x2xi8>, tensor<2xf32>, tensor<2xi8>)
      outs(%bi : tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x2x2x2xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  %result = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>}
      ins(%a, %b : tensor<1x4x5x2xf32>, tensor<2x2x2xf32>)
      outs(%init : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  check.expect_eq_const(%result, dense<[[[[52.0, -5.0], [-7.0, -2.0]], [[-1.0, -1.5], [35.0, 30.0]]]]> : tensor<1x2x2x2xf32>) : tensor<1x2x2x2xf32>
  return
}

// A custom strided window contraction with interleaved reduction/parallel
// loops, a broadcast RHS batch, and output axes ordered (n, batch, m).
// C[n,t,m] = sum_{r,c} DQ(A[t,2*m+r,c]) * DQ(B[n,c,r]).
func.func @generic_strided_permuted_contraction() {
  %aq = util.unfoldable_constant dense<[[[128, 131], [132, 135], [136, 139], [140, 143], [144, 147], [148, 128], [129, 132]], [[137, 140], [141, 144], [145, 148], [149, 129], [130, 133], [134, 137], [138, 141]]]> : tensor<2x7x2xi8>
  %bq = util.unfoldable_constant dense<[[[-6, -4, -2], [-3, -1, 1]], [[-1, 1, 3], [2, 4, 6]], [[4, 6, -5], [-6, -4, -2]]]> : tensor<3x2x3xi8>
  %sa = util.unfoldable_constant dense<[0.5, 1.0]> : tensor<2xf32>
  %za = util.unfoldable_constant dense<[130, 134]> : tensor<2xi8>
  %sb = util.unfoldable_constant dense<[1.0, 0.5, 2.0]> : tensor<3xf32>
  %zb = util.unfoldable_constant dense<[-2, 1, 3]> : tensor<3xi8>
  %ai = tensor.empty() : tensor<2x7x2xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, zp_unsigned, indexing_maps = [affine_map<(t, x, c) -> (t, x, c)>,
                        affine_map<(t, x, c) -> (t)>,
                        affine_map<(t, x, c) -> (t)>,
                        affine_map<(t, x, c) -> (t, x, c)>]>
      ins(%aq, %sa, %za : tensor<2x7x2xi8>, tensor<2xf32>, tensor<2xi8>)
      outs(%ai : tensor<2x7x2xf32>) -> tensor<2x7x2xf32>
  %bi = tensor.empty() : tensor<3x2x3xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(n, c, r) -> (n, c, r)>,
                        affine_map<(n, c, r) -> (n)>,
                        affine_map<(n, c, r) -> (n)>,
                        affine_map<(n, c, r) -> (n, c, r)>]>
      ins(%bq, %sb, %zb : tensor<3x2x3xi8>, tensor<3xf32>, tensor<3xi8>)
      outs(%bi : tensor<3x2x3xf32>) -> tensor<3x2x3xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<3x2x3xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<3x2x3xf32>) -> tensor<3x2x3xf32>
  %result = linalg.generic {
      indexing_maps = [affine_map<(c, n, t, r, m) -> (t, m * 2 + r, c)>,
                       affine_map<(c, n, t, r, m) -> (n, c, r)>,
                       affine_map<(c, n, t, r, m) -> (n, t, m)>],
      iterator_types = ["reduction", "parallel", "parallel", "reduction", "parallel"]}
      ins(%a, %b : tensor<2x7x2xf32>, tensor<3x2x3xf32>)
      outs(%init : tensor<3x2x3xf32>) {
    ^bb0(%lhs: f32, %rhs: f32, %acc: f32):
      %product = arith.mulf %lhs, %rhs : f32
      %sum = arith.addf %acc, %product : f32
      linalg.yield %sum : f32
  } -> tensor<3x2x3xf32>
  check.expect_eq_const(%result, dense<[[[17.5, 5.5, -52.5], [20.0, -96.0, 41.0]], [[19.25, 37.25, -2.25], [61.0, -18.0, 29.5]], [[-133.0, -333.0, -73.0], [-516.0, 4.0, -166.0]]]> : tensor<3x2x3xf32>) : tensor<3x2x3xf32>
  return
}

// Both asymmetric operands vary by block and output axis, with different
// storage permutations folded into dequantization. Only k is reduced in i32;
// g is retained in partials and reduced after applying both block scales.
// C[n,m] = sum_{g,k} DQ(A[k,m,g]) * DQ(B[n,k,g]).
func.func @blockwise_asymmetric_permuted_storage() {
  %aq = util.unfoldable_constant dense<[[[128, 131], [133, 136]], [[135, 138], [140, 143]], [[142, 145], [147, 129]]]> : tensor<3x2x2xi8>
  %bq = util.unfoldable_constant dense<[[[-8, -3], [-5, 0], [-2, 3]], [[-1, 4], [2, 7], [5, -7]], [[6, -6], [-8, -3], [-5, 0]]]> : tensor<3x3x2xi8>
  %sa = util.unfoldable_constant dense<[[0.5, 2.0], [1.0, 0.25]]> : tensor<2x2xf32>
  %za = util.unfoldable_constant dense<[[130, 135], [134, 129]]> : tensor<2x2xi8>
  %sb = util.unfoldable_constant dense<[[1.0, 0.5], [2.0, 0.25], [0.5, 2.0]]> : tensor<3x2xf32>
  %zb = util.unfoldable_constant dense<[[-3, 1], [2, -4], [-1, 3]]> : tensor<3x2xi8>
  %ai = tensor.empty() : tensor<2x2x3xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, zp_unsigned, indexing_maps = [affine_map<(m, g, k) -> (k, m, g)>,
                        affine_map<(m, g, k) -> (g, m)>,
                        affine_map<(m, g, k) -> (g, m)>,
                        affine_map<(m, g, k) -> (m, g, k)>]>
      ins(%aq, %sa, %za : tensor<3x2x2xi8>, tensor<2x2xf32>, tensor<2x2xi8>)
      outs(%ai : tensor<2x2x3xf32>) -> tensor<2x2x3xf32>
  %bi = tensor.empty() : tensor<2x3x3xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(g, k, n) -> (n, k, g)>,
                        affine_map<(g, k, n) -> (n, g)>,
                        affine_map<(g, k, n) -> (n, g)>,
                        affine_map<(g, k, n) -> (g, k, n)>]>
      ins(%bq, %sb, %zb : tensor<3x3x2xi8>, tensor<3x2xf32>, tensor<3x2xi8>)
      outs(%bi : tensor<2x3x3xf32>) -> tensor<2x3x3xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<3x2xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<3x2xf32>) -> tensor<3x2xf32>
  %result = linalg.contract
      indexing_maps = [affine_map<(m, n, g, k) -> (m, g, k)>,
                       affine_map<(m, n, g, k) -> (g, k, n)>,
                       affine_map<(m, n, g, k) -> (n, m)>]
      ins(%a, %b : tensor<2x2x3xf32>, tensor<2x3x3xf32>)
      outs(%init : tensor<3x2xf32>) -> tensor<3x2xf32>
  check.expect_eq_const(%result, dense<[[21.0, 18.75], [38.75, 181.125], [-84.25, -170.5]]> : tensor<3x2xf32>) : tensor<3x2xf32>
  return
}

// Anticipated QDQ Linalg translation of a PT2E/XNNPACK bias-free Linear(16,16)
// with an 8x16 input: Q -> DQ -> matmul(input, DQ(weight).T) -> Q -> DQ.
// Input, quantized weights, scales, zero points, and expected output are from
// a seed-0 PyTorch run. The RHS indexing map incorporates aten.transpose.int.
func.func @pt2e_linear_qdq() {
  %input_scale = arith.constant 0.024613387882709503 : f32
  %weight_scale = arith.constant 0.0019636377692222595 : f32
  %output_scale = arith.constant 0.010716261342167854 : f32
  %input_zp = arith.constant -33 : i64
  %weight_zp = arith.constant 0 : i64
  %output_zp = arith.constant -10 : i64
  %zero = arith.constant 0.0 : f32

  // Prevent constant folding so the QDQ graph executes at runtime.
  %input = util.unfoldable_constant dense<[
      [-0.03937254101037979, -0.8014721870422363, -0.49554431438446045, -0.36151406168937683, 0.5851132273674011, -1.1560065746307373, -0.14336487650871277, -0.19474059343338013, -0.08556340634822845, 1.3945199251174927, 0.5969000458717346, -0.4828483462333679, -0.3660986125469208, -1.3270524740219116, 1.695279598236084, 2.0654995441436768],
      [-0.2339576780796051, 0.7073183059692383, 0.5800480842590332, 0.2683020830154419, -2.05893874168396, 0.5340206623077393, -0.5353949069976807, -0.863664984703064, -0.023494405671954155, 1.171669840812683, 0.3986871838569641, -0.19871552288532257, -1.1559407711029053, -0.3166695237159729, 0.9402980804443359, -1.1469545364379883],
      [0.5588033199310303, 0.7917584776878357, -0.18467576801776886, -0.7317724823951721, -0.0806519165635109, -0.9800607562065125, 0.060491468757390976, -0.4889546036720276, -0.8137312531471252, 0.8199948072433472, -0.633173406124115, 1.2947547435760498, 1.4628292322158813, -0.6204336285591125, 0.988389253616333, -0.4321781396865845],
      [-0.6232188940048218, -0.21624533832073212, -0.48867544531822205, 0.7869559526443481, 0.10759072750806808, -1.0714776515960693, -0.11664776504039764, -1.0169707536697388, -1.1979783773422241, 0.47843775153160095, -1.2295243740081787, -1.3700363636016846, 1.5435417890548706, -0.0332069993019104, -0.4186263680458069, -0.25559672713279724],
      [-0.12923042476177216, -0.05459475517272949, 0.4083467423915863, 1.1263659000396729, 1.9350574016571045, 1.0076850652694702, 1.0046416521072388, -0.43351978063583374, -1.2425976991653442, 1.2845500707626343, 0.24377228319644928, 0.5303686261177063, -0.014530729502439499, -2.235718250274658, 1.4660301208496094, -1.2190581560134888],
      [0.64423006772995, 3.9300038814544678, -0.1244242936372757, 0.29534170031547546, 0.38265419006347656, -0.549721360206604, -0.9940357208251953, 1.345936894416809, 1.9456682205200195, -1.2903639078140259, -2.3494760990142822, -2.068861961364746, 0.9094210863113403, -0.6946200728416443, 1.9594571590423584, -1.1038278341293335],
      [0.5411418080329895, 1.5389583110809326, 1.0860490798950195, 1.246405005455017, 0.11507505923509598, 1.619307518005371, 0.4636935293674469, 1.3007354736328125, 0.873230516910553, 0.0651267021894455, 0.7732411026954651, -0.970138430595398, -0.8876768946647644, -0.3183162808418274, -0.3344036936759949, 0.4542836546897888],
      [0.49895304441452026, 0.8779974579811096, 0.38944435119628906, 1.462517499923706, 0.47950610518455505, -0.5333999991416931, -0.03465134650468826, 0.6572969555854797, -0.3112243115901947, -0.5620035529136658, -0.4834926426410675, -1.2721126079559326, -0.17401821911334991, 0.5541167259216309, -0.18165524303913116, -0.23447339236736298]
  ]> : tensor<8x16xf32>
  // Signed bytes decoded from torch_tensor_16_16_torch.int8, excluding its
  // four-byte resource alignment header. Rows retain Torch weight layout [N,K].
  %weight = util.unfoldable_constant dense<[
      [-1, 68, -105, -94, -49, 34, -3, 101, -11, 34, -38, -25, -122, -84, -52, 5],
      [50, 76, -86, -55, 46, 106, -26, 95, -21, 13, 115, -118, -80, -32, -50, 110],
      [-83, -59, -89, -119, -74, 109, 57, 62, 7, -65, 22, -119, -92, -66, 80, 75],
      [-56, -5, 81, 127, 51, 17, 85, -75, 24, -99, -88, -66, 58, 51, -75, 38],
      [70, -16, 5, 29, 79, 122, -98, -47, 50, 105, 111, 112, 25, -111, 12, -80],
      [-119, 113, 97, -127, 24, -21, -21, -58, 49, -75, 47, 64, 91, 48, -126, -83],
      [64, 27, -99, -73, 120, 86, -56, -32, -121, -2, -96, -98, -7, 19, -52, 76],
      [-77, 116, 87, -107, -32, 6, 19, 30, 50, 8, -62, 60, -122, -75, -32, -62],
      [-45, -104, -27, 27, -83, -7, 91, -13, 4, -11, 26, 81, 121, 81, 121, -9],
      [-114, -60, 87, -1, -63, -98, -119, -107, -26, 70, 69, -123, 79, -100, -27, -52],
      [-25, -25, -114, -110, -20, 2, -58, 48, -115, -9, 112, -52, 115, 46, -115, 81],
      [-15, -57, 102, -103, 14, -27, 91, 36, 61, 45, -31, -27, -105, 69, 101, 87],
      [-90, 6, -90, -70, -74, 44, -76, -3, 5, 82, -96, -87, -74, 89, -46, 107],
      [46, 16, -1, -25, 16, -29, -1, 16, -100, -67, 103, -103, -9, 126, 46, 4],
      [-110, 63, -91, -36, -43, -19, 1, 105, 16, 114, 78, -80, 57, -90, -54, 37],
      [42, 96, -41, 0, 66, -123, 92, -105, 2, -22, -67, 17, 105, -37, -76, -47]
  ]> : tensor<16x16xi8>
  %input_empty = tensor.empty() : tensor<8x16xf32>

  %input_q_empty = tensor.empty() : tensor<8x16xi8>
  %input_q = iree_linalg_ext.quantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>, affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>],
       quant_min = -128, quant_max = 127>
      ins(%input, %input_scale, %input_zp : tensor<8x16xf32>, f32, i64)
      outs(%input_q_empty : tensor<8x16xi8>) -> tensor<8x16xi8>
  %input_dq = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>, affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%input_q, %input_scale, %input_zp : tensor<8x16xi8>, f32, i64)
      outs(%input_empty : tensor<8x16xf32>) -> tensor<8x16xf32>
  %weight_dq_empty = tensor.empty() : tensor<16x16xf32>
  %weight_dq = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(n, k) -> (n, k)>,
                        affine_map<(n, k) -> ()>, affine_map<(n, k) -> ()>,
                        affine_map<(n, k) -> (n, k)>]>
      ins(%weight, %weight_scale, %weight_zp : tensor<16x16xi8>, f32, i64)
      outs(%weight_dq_empty : tensor<16x16xf32>) -> tensor<16x16xf32>
  %init = linalg.fill ins(%zero : f32) outs(%input_empty : tensor<8x16xf32>) -> tensor<8x16xf32>
  %linear = linalg.contract
      indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
                       affine_map<(m, n, k) -> (n, k)>,
                       affine_map<(m, n, k) -> (m, n)>]
      ins(%input_dq, %weight_dq : tensor<8x16xf32>, tensor<16x16xf32>)
      outs(%init : tensor<8x16xf32>) -> tensor<8x16xf32>
  %output_q = iree_linalg_ext.quantize_affine
      <indexing_maps = [affine_map<(m, n) -> (m, n)>,
                        affine_map<(m, n) -> ()>, affine_map<(m, n) -> ()>,
                        affine_map<(m, n) -> (m, n)>],
       quant_min = -128, quant_max = 127>
      ins(%linear, %output_scale, %output_zp : tensor<8x16xf32>, f32, i64)
      outs(%input_q_empty : tensor<8x16xi8>) -> tensor<8x16xi8>
  %output = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, n) -> (m, n)>,
                        affine_map<(m, n) -> ()>, affine_map<(m, n) -> ()>,
                        affine_map<(m, n) -> (m, n)>]>
      ins(%output_q, %output_scale, %output_zp : tensor<8x16xi8>, f32, i64)
      outs(%input_empty : tensor<8x16xf32>) -> tensor<8x16xf32>

  // Recorded torch_quantized_output. The original IREE lowering matched
  // PyTorch exactly; require the same result with and without the QDQ rewrite.
  check.expect_eq_const(%output, dense<[
      [0.11787887662649155, 0.4929480254650116, 0.67512446641922, -0.7179895043373108, 0.13931140303611755, -1.2323700189590454, 0.16074392199516296, -0.43936672806739807, 0.06429757177829742, 0.5465293526649475, 0.11787887662649155, 0.7930033206939697, 0.2786228060722351, -0.03214878588914871, 0.6215431690216064, -0.3214878439903259],
      [0.2786228060722351, -0.2250414937734604, 0.18217644095420837, -0.6322594285011292, 0.4607992470264435, -0.18217644095420837, -0.7287057638168335, 0.6644082069396973, -0.010716261342167854, 0.7501382827758789, -0.7394220232963562, 0.02143252268433571, 0.2786228060722351, -0.15002766251564026, 0.09644635021686554, -0.6108269095420837],
      [-0.17146018147468567, -0.7930033206939697, -0.7715708017349243, -0.5786781311035156, 0.33220410346984863, 0.2143252193927765, -0.05358130484819412, 0.08573009073734283, 0.4500829875469208, 0.12859514355659485, -0.04286504536867142, -0.2893390655517578, -0.5143805742263794, -0.2786228060722351, -0.010716261342167854, 0.8787334561347961],
      [-0.4607992470264435, -0.6429756879806519, -0.5358130931854248, 0.7501382827758789, -0.6429756879806519, -0.07501383125782013, 0.5786781311035156, -0.7287057638168335, 0.13931140303611755, 1.103774905204773, 0.428650438785553, -0.6536919474601746, 0.2893390655517578, 0.06429757177829742, 0.2571902871131897, 0.9001659750938416],
      [-0.2357577532529831, -0.2143252193927765, -0.2679065465927124, -0.12859514355659485, 1.4681278467178345, -0.7179895043373108, 0.0, 0.06429757177829742, -0.010716261342167854, 0.15002766251564026, -1.0394773483276367, -0.2786228060722351, -1.2645188570022583, -0.4929480254650116, -0.07501383125782013, 0.20360895991325378],
      [0.5250968337059021, 0.17146018147468567, -0.08573009073734283, 0.2786228060722351, -0.6429756879806519, 0.2464740127325058, 0.2679065465927124, 0.8358683586120605, -0.921598494052887, -0.43936672806739807, -1.2323700189590454, -0.2893390655517578, -0.10716260969638824, -0.03214878588914871, 0.2893390655517578, 0.8787334561347961],
      [0.38578540086746216, 1.1895049810409546, 0.2679065465927124, 0.33220410346984863, 0.3107715845108032, -0.2679065465927124, -0.12859514355659485, 0.40721791982650757, -0.8037195801734924, -0.6322594285011292, -0.5572456121444702, 0.12859514355659485, -0.16074392199516296, 0.06429757177829742, 0.33220410346984863, -0.5893943905830383],
      [-0.15002766251564026, 0.15002766251564026, -0.4607992470264435, 0.67512446641922, -0.6215431690216064, -0.2893390655517578, 0.2250414937734604, -0.2357577532529831, -0.4607992470264435, -0.16074392199516296, -0.3107715845108032, -0.20360895991325378, -0.16074392199516296, 0.4929480254650116, -0.18217644095420837, 0.2464740127325058]
  ]> : tensor<8x16xf32>) : tensor<8x16xf32>
  return
}

// Execution counterparts of numerical corner cases in
// GlobalOptimization/test/convert_qdq_to_integer_math_scale_range.mlir.
// Tests ending in _reference or _rewritten document accepted numerical
// differences. The build targets use --gtest_filter to select their respective
// expectations; tests without these suffixes run in both configurations.
// Unfolded storage operands prevent constant folding of the contractions.

// (1 * 1e20) * (0 * 1e20) = 0. Sequential scaling avoids the overflowing
// scale product that previously caused the integer rewrite to return NaN.
func.func @scale_product_overflow() {
  %aq = util.unfoldable_constant dense<1> : tensor<1x1xi8>
  %bq = util.unfoldable_constant dense<0> : tensor<1x1xi8>
  %scale = arith.constant 1.0e20 : f32
  %ai = tensor.empty() : tensor<1x1xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %scale : tensor<1x1xi8>, f32)
      outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32>
  %bi = tensor.empty() : tensor<1x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %scale : tensor<1x1xi8>, f32)
      outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  check.expect_eq_const(%result, dense<0.0> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

// f16 dequantization rounds 2049 to 2048, so -2048 + 2048 = 0. The integer
// rewrite skips that rounding and returns 1.
func.func private @compute_f16_dequantization_rounding() -> tensor<1x1xf16> {
  %aq = util.unfoldable_constant dense<[[-1, 1]]> : tensor<1x2xi8>
  %bq = util.unfoldable_constant dense<[[2048], [2049]]> : tensor<2x1xi16>
  %scale = arith.constant 1.0 : f32
  %ai = tensor.empty() : tensor<1x2xf16>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %scale : tensor<1x2xi8>, f32)
      outs(%ai : tensor<1x2xf16>) -> tensor<1x2xf16>
  %bi = tensor.empty() : tensor<2x1xf16>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %scale : tensor<2x1xi16>, f32)
      outs(%bi : tensor<2x1xf16>) -> tensor<2x1xf16>
  %zero = arith.constant 0.0 : f16
  %empty = tensor.empty() : tensor<1x1xf16>
  %init = linalg.fill ins(%zero : f16) outs(%empty : tensor<1x1xf16>) -> tensor<1x1xf16>
  %result = linalg.matmul ins(%a, %b : tensor<1x2xf16>, tensor<2x1xf16>)
      outs(%init : tensor<1x1xf16>) -> tensor<1x1xf16>
  return %result : tensor<1x1xf16>
}

func.func @f16_dequantization_rounding_reference() {
  %result = call @compute_f16_dequantization_rounding() : () -> tensor<1x1xf16>
  check.expect_eq_const(%result, dense<0.0> : tensor<1x1xf16>) : tensor<1x1xf16>
  return
}

func.func @f16_dequantization_rounding_rewritten() {
  %result = call @compute_f16_dequantization_rounding() : () -> tensor<1x1xf16>
  check.expect_eq_const(%result, dense<1.0> : tensor<1x1xf16>) : tensor<1x1xf16>
  return
}

// -0 + (0 * -1) = -0. The integer rewrite discards the init sign and
// returns +0.
func.func private @compute_negative_zero_init() -> tensor<1x1xf32> {
  %aq = util.unfoldable_constant dense<0> : tensor<1x1xi8>
  %bq = util.unfoldable_constant dense<-1> : tensor<1x1xi8>
  %scale = arith.constant 1.0 : f32
  %ai = tensor.empty() : tensor<1x1xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %scale : tensor<1x1xi8>, f32)
      outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32>
  %bi = tensor.empty() : tensor<1x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %scale : tensor<1x1xi8>, f32)
      outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32>
  %zero = arith.constant -0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  return %result : tensor<1x1xf32>
}

func.func @negative_zero_init_reference() {
  %result = call @compute_negative_zero_init() : () -> tensor<1x1xf32>
  // Floating-point equality cannot distinguish -0 from +0; check the bits.
  %c0 = arith.constant 0 : index
  %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32>
  %bits = arith.bitcast %value : f32 to i32
  %negative_zero_bits = arith.constant -2147483648 : i32
  %is_negative_zero = arith.cmpi eq, %bits, %negative_zero_bits : i32
  check.expect_true(%is_negative_zero) : i1
  return
}

func.func @negative_zero_init_rewritten() {
  %result = call @compute_negative_zero_init() : () -> tensor<1x1xf32>
  // Floating-point equality cannot distinguish -0 from +0; check the bits.
  %c0 = arith.constant 0 : index
  %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32>
  %bits = arith.bitcast %value : f32 to i32
  %positive_zero_bits = arith.constant 0 : i32
  %is_positive_zero = arith.cmpi eq, %bits, %positive_zero_bits : i32
  check.expect_true(%is_positive_zero) : i1
  return
}

// [0 * infinity, 1 * infinity] dot [1, 1] is NaN. The integer rewrite
// instead computes 1 * infinity and returns infinity.
func.func private @compute_infinite_scale_nan_propagation() -> tensor<1x1xf32> {
  %aq = util.unfoldable_constant dense<[[0, 1]]> : tensor<1x2xi8>
  %bq = util.unfoldable_constant dense<1> : tensor<2x1xi8>
  %sa = arith.constant 0x7F800000 : f32
  %sb = arith.constant 1.0 : f32
  %ai = tensor.empty() : tensor<1x2xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %sa : tensor<1x2xi8>, f32)
      outs(%ai : tensor<1x2xf32>) -> tensor<1x2xf32>
  %bi = tensor.empty() : tensor<2x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %sb : tensor<2x1xi8>, f32)
      outs(%bi : tensor<2x1xf32>) -> tensor<2x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x2xf32>, tensor<2x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  return %result : tensor<1x1xf32>
}

func.func @infinite_scale_nan_propagation_reference() {
  %result = call @compute_infinite_scale_nan_propagation() : () -> tensor<1x1xf32>
  // Compare unordered rather than requiring a particular NaN payload.
  %c0 = arith.constant 0 : index
  %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32>
  %is_nan = arith.cmpf uno, %value, %value : f32
  check.expect_true(%is_nan) : i1
  return
}

func.func @infinite_scale_nan_propagation_rewritten() {
  %result = call @compute_infinite_scale_nan_propagation() : () -> tensor<1x1xf32>
  check.expect_eq_const(%result, dense<0x7F800000> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

// (32767 * 1e-23)^2 is approximately 1.07368e-37. Sequential scaling avoids
// the underflowing scale product that previously erased the result.
func.func @scale_product_underflow() {
  %q = util.unfoldable_constant dense<32767> : tensor<1x1xi16>
  %scale = arith.constant 1.0e-23 : f32
  %ai = tensor.empty() : tensor<1x1xf32>
  %a = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%q, %scale : tensor<1x1xi16>, f32)
      outs(%ai : tensor<1x1xf32>) -> tensor<1x1xf32>
  %bi = tensor.empty() : tensor<1x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%q, %scale : tensor<1x1xi16>, f32)
      outs(%bi : tensor<1x1xf32>) -> tensor<1x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x1xf32>, tensor<1x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  // The dequantized operands and their product are all normal f32 values.
  // Use exact equality: an absolute tolerance could incorrectly accept zero.
  check.expect_eq_const(%result, dense<0x02122428> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

// Both reference operands are finite, but D * 1e30 overflows in the epilogue.
func.func private @compute_opposing_scales_large_lhs() -> tensor<1x1xf32> {
  %aq = util.unfoldable_constant dense<-1> : tensor<1x32768xi8>
  %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8>
  %sa = arith.constant 1.0e30 : f32
  %sb = arith.constant 1.0e-30 : f32
  %ai = tensor.empty() : tensor<1x32768xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %sa : tensor<1x32768xi8>, f32)
      outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32>
  %bi = tensor.empty() : tensor<32768x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %sb : tensor<32768x1xi8>, f32)
      outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  return %result : tensor<1x1xf32>
}

func.func @opposing_scales_large_lhs_reference() {
  %result = call @compute_opposing_scales_large_lhs() : () -> tensor<1x1xf32>
  // D = 32768 * 255 * 127 = 1061191680, within signed i32.
  // Allow reduction-order rounding while requiring a finite result.
  check.expect_almost_eq_const(%result, dense<1061191680.0> : tensor<1x1xf32>, atol 0.0, rtol 1.0e-4) : tensor<1x1xf32>
  return
}

func.func @opposing_scales_large_lhs_rewritten() {
  %result = call @compute_opposing_scales_large_lhs() : () -> tensor<1x1xf32>
  // 0x7F800000 -> INFINITY
  check.expect_eq_const(%result, dense<0x7F800000> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

// Reversing the scales avoids epilogue overflow. Floating-point reduction
// rounding differs; allow 1e-4 relative error against the exact integer sum.
func.func private @compute_opposing_scales_large_rhs() -> tensor<1x1xf32> {
  %aq = util.unfoldable_constant dense<-1> : tensor<1x32768xi8>
  %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8>
  %sa = arith.constant 1.0e-30 : f32
  %sb = arith.constant 1.0e30 : f32
  %ai = tensor.empty() : tensor<1x32768xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %sa : tensor<1x32768xi8>, f32)
      outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32>
  %bi = tensor.empty() : tensor<32768x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %sb : tensor<32768x1xi8>, f32)
      outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  return %result : tensor<1x1xf32>
}

func.func @opposing_scales_large_rhs_reference() {
  %result = call @compute_opposing_scales_large_rhs() : () -> tensor<1x1xf32>
  // D = 32768 * 255 * 127 = 1061191680, within signed i32.
  // Allow reduction-order rounding while requiring a finite result.
  check.expect_almost_eq_const(%result, dense<1061191680.0> : tensor<1x1xf32>, atol 0.0, rtol 1.0e-4) : tensor<1x1xf32>
  return
}

func.func @opposing_scales_large_rhs_rewritten() {
  %result = call @compute_opposing_scales_large_rhs() : () -> tensor<1x1xf32>
  check.expect_eq_const(%result, dense<1061191680.0> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

// All scales are finite, but 127 * 1e37 overflows during reference
// dequantization. Integer reduction bypasses that overflow.
func.func private @compute_finite_scale_dequantization_overflow() -> tensor<1x1xf32> {
  %aq = util.unfoldable_constant dense<-1> : tensor<1x32768xi8>
  %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8>
  %sa = arith.constant 1.0e-37 : f32
  %sb = arith.constant 1.0e37 : f32
  %ai = tensor.empty() : tensor<1x32768xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %sa : tensor<1x32768xi8>, f32)
      outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32>
  %bi = tensor.empty() : tensor<32768x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %sb : tensor<32768x1xi8>, f32)
      outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  return %result : tensor<1x1xf32>
}

func.func @finite_scale_dequantization_overflow_reference() {
  %result = call @compute_finite_scale_dequantization_overflow() : () -> tensor<1x1xf32>
  check.expect_eq_const(%result, dense<0x7F800000> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

func.func @finite_scale_dequantization_overflow_rewritten() {
  %result = call @compute_finite_scale_dequantization_overflow() : () -> tensor<1x1xf32>
  check.expect_eq_const(%result, dense<1061191616.0> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}

// With zero activations, the same overflowing reference weights yield NaN
// from 0 * infinity. Integer reduction instead yields zero.
func.func private @compute_finite_scale_dequantization_nan() -> tensor<1x1xf32> {
  %aq = util.unfoldable_constant dense<0> : tensor<1x32768xi8>
  %bq = util.unfoldable_constant dense<127> : tensor<32768x1xi8>
  %sa = arith.constant 1.0e-37 : f32
  %sb = arith.constant 1.0e37 : f32
  %ai = tensor.empty() : tensor<1x32768xf32>
  %a = iree_linalg_ext.dequantize_affine
      <input_unsigned, indexing_maps = [affine_map<(m, k) -> (m, k)>,
                        affine_map<(m, k) -> ()>,
                        affine_map<(m, k) -> (m, k)>]>
      ins(%aq, %sa : tensor<1x32768xi8>, f32)
      outs(%ai : tensor<1x32768xf32>) -> tensor<1x32768xf32>
  %bi = tensor.empty() : tensor<32768x1xf32>
  %b = iree_linalg_ext.dequantize_affine
      <indexing_maps = [affine_map<(k, n) -> (k, n)>,
                        affine_map<(k, n) -> ()>,
                        affine_map<(k, n) -> (k, n)>]>
      ins(%bq, %sb : tensor<32768x1xi8>, f32)
      outs(%bi : tensor<32768x1xf32>) -> tensor<32768x1xf32>
  %zero = arith.constant 0.0 : f32
  %empty = tensor.empty() : tensor<1x1xf32>
  %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x1xf32>) -> tensor<1x1xf32>
  %result = linalg.matmul ins(%a, %b : tensor<1x32768xf32>, tensor<32768x1xf32>)
      outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32>
  return %result : tensor<1x1xf32>
}

func.func @finite_scale_dequantization_nan_reference() {
  %result = call @compute_finite_scale_dequantization_nan() : () -> tensor<1x1xf32>
  // Compare unordered without requiring a specific NaN payload.
  %c0 = arith.constant 0 : index
  %value = tensor.extract %result[%c0, %c0] : tensor<1x1xf32>
  %is_nan = arith.cmpf uno, %value, %value : f32
  check.expect_true(%is_nan) : i1
  return
}

func.func @finite_scale_dequantization_nan_rewritten() {
  %result = call @compute_finite_scale_dequantization_nan() : () -> tensor<1x1xf32>
  check.expect_eq_const(%result, dense<0.0> : tensor<1x1xf32>) : tensor<1x1xf32>
  return
}