// Test large aligned linalg matmul to make sure we go through the optimized
// path for GPUs.

// Problem size      : 2048x512x1024
// Input type        : F32
// Accumulation type : F32
func.func @matmul_2048x512x1024_f32_f32() {
  %lhs = util.unfoldable_constant dense<1.0> : tensor<2048x1024xf32>
  %rhs = util.unfoldable_constant dense<0.4> : tensor<1024x512xf32>
  %c0 = arith.constant 0.0 : f32
  %init = linalg.init_tensor[2048, 512] : tensor<2048x512xf32>
  %CC = linalg.fill ins(%c0 : f32) outs(%init : tensor<2048x512xf32>) -> tensor<2048x512xf32>
  %D = linalg.matmul ins(%lhs, %rhs: tensor<2048x1024xf32>, tensor<1024x512xf32>)
                    outs(%CC: tensor<2048x512xf32>) -> tensor<2048x512xf32>
  check.expect_almost_eq_const(%D, dense<409.596> : tensor<2048x512xf32>) : tensor<2048x512xf32>
  return
}

// Problem size      : 3456x1024x2048
// Input type        : F16
// Accumulation type : F16 
func.func @matmul_3456x1024x2048_f16_f16() {
  %lhs = util.unfoldable_constant dense<1.00> : tensor<3456x2048xf16>
  %rhs = util.unfoldable_constant dense<0.01> : tensor<2048x1024xf16>
  %c0 = arith.constant 0.0 : f16
  %init = linalg.init_tensor[3456, 1024] : tensor<3456x1024xf16>
  %CC = linalg.fill ins(%c0 : f16) outs(%init : tensor<3456x1024xf16>) -> tensor<3456x1024xf16>
  %D = linalg.matmul ins(%lhs, %rhs: tensor<3456x2048xf16>, tensor<2048x1024xf16>)
                    outs(%CC: tensor<3456x1024xf16>) -> tensor<3456x1024xf16>
  check.expect_almost_eq_const(%D, dense<20.2812> : tensor<3456x1024xf16>) : tensor<3456x1024xf16>
  return
}