| // End-to-end correctness test for the batch-matmul-to-matmul pass |
| // (--iree-global-opt-convert-batch-matmul-to-matmul), which runs as |
| // part of the default global optimization pipeline. |
| |
| func.func @broadcast_batch_matmul() { |
| %act = util.unfoldable_constant dense<[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], |
| [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]]> : tensor<2x2x3xf32> |
| %weight = util.unfoldable_constant dense<[[1.0, 0.0, 0.0, 1.0], |
| [0.0, 1.0, 0.0, 1.0], |
| [0.0, 0.0, 1.0, 1.0]]> : tensor<3x4xf32> |
| %cst = arith.constant 0.0 : f32 |
| %init_broadcast = tensor.empty() : tensor<2x3x4xf32> |
| %broadcast = linalg.broadcast ins(%weight : tensor<3x4xf32>) |
| outs(%init_broadcast : tensor<2x3x4xf32>) dimensions = [0] |
| %init_out = tensor.empty() : tensor<2x2x4xf32> |
| %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<2x2x4xf32>) -> tensor<2x2x4xf32> |
| |
| %opt = linalg.batch_matmul ins(%act, %broadcast : tensor<2x2x3xf32>, tensor<2x3x4xf32>) |
| outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32> |
| |
| %broadcast_ref = util.optimization_barrier %broadcast : tensor<2x3x4xf32> |
| %ref = linalg.batch_matmul ins(%act, %broadcast_ref : tensor<2x2x3xf32>, tensor<2x3x4xf32>) |
| outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32> |
| |
| check.expect_almost_eq(%opt, %ref) : tensor<2x2x4xf32> |
| return |
| } |
| |
| func.func @broadcast_batch_matmul_transpose_b() { |
| %act = util.unfoldable_constant dense<[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], |
| [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]]> : tensor<2x2x3xf32> |
| %weight = util.unfoldable_constant dense<[[1.0, 0.0, 0.0], |
| [0.0, 1.0, 0.0], |
| [0.0, 0.0, 1.0], |
| [1.0, 1.0, 1.0]]> : tensor<4x3xf32> |
| %cst = arith.constant 0.0 : f32 |
| %init_broadcast = tensor.empty() : tensor<2x4x3xf32> |
| %broadcast = linalg.broadcast ins(%weight : tensor<4x3xf32>) |
| outs(%init_broadcast : tensor<2x4x3xf32>) dimensions = [0] |
| %init_out = tensor.empty() : tensor<2x2x4xf32> |
| %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<2x2x4xf32>) -> tensor<2x2x4xf32> |
| |
| %opt = linalg.batch_matmul |
| indexing_maps = [affine_map<(b, m, n, k) -> (b, m, k)>, |
| affine_map<(b, m, n, k) -> (b, n, k)>, |
| affine_map<(b, m, n, k) -> (b, m, n)>] |
| ins(%act, %broadcast : tensor<2x2x3xf32>, tensor<2x4x3xf32>) |
| outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32> |
| |
| %broadcast_ref = util.optimization_barrier %broadcast : tensor<2x4x3xf32> |
| %ref = linalg.batch_matmul |
| indexing_maps = [affine_map<(b, m, n, k) -> (b, m, k)>, |
| affine_map<(b, m, n, k) -> (b, n, k)>, |
| affine_map<(b, m, n, k) -> (b, m, n)>] |
| ins(%act, %broadcast_ref : tensor<2x2x3xf32>, tensor<2x4x3xf32>) |
| outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32> |
| |
| check.expect_almost_eq(%opt, %ref) : tensor<2x2x4xf32> |
| return |
| } |
