| // This file contains tests of linalg dialect operations with dynamic shapes. |
| |
| #map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> |
| #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)> |
| #map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> |
| |
| // Batched matmul where B=32 M=1 N=128 K=dynamic. |
| // Tested with K=3 and splat operands. |
| func.func @batch_matmul_dynamic_reduction_size_B32_M1_N128_K3() { |
| %lhs = flow.tensor.dynamic_constant dense<1.0> : tensor<32x1x3xf16> -> tensor<32x1x?xf16> |
| %rhs = flow.tensor.dynamic_constant dense<2.0> : tensor<32x3x128xf16> -> tensor<32x?x128xf16> |
| |
| %cst = arith.constant 0.000000 : f16 |
| %2 = tensor.empty() : tensor<32x1x128xf16> |
| %3 = linalg.fill ins(%cst : f16) outs(%2 : tensor<32x1x128xf16>) -> tensor<32x1x128xf16> |
| %observed = linalg.generic { |
| indexing_maps = [#map, #map1, #map2], |
| iterator_types = ["parallel", "parallel", "parallel", "reduction"]} |
| ins(%lhs, %rhs : tensor<32x1x?xf16>, tensor<32x?x128xf16>) outs(%3 : tensor<32x1x128xf16>) { |
| ^bb0(%in: f16, %in_0: f16, %out: f16): |
| %6 = arith.mulf %in, %in_0 : f16 |
| %7 = arith.addf %out, %6 : f16 |
| linalg.yield %7 : f16 |
| } -> tensor<32x1x128xf16> |
| |
| %expected = flow.tensor.dynamic_constant dense<6.0> : tensor<32x1x128xf16> -> tensor<32x1x128xf16> |
| check.expect_almost_eq(%observed, %expected, atol 1.0e-04) : tensor<32x1x128xf16> |
| return |
| } |
| |
| // Batched matmul where B=2 M=1 N=3 K=dynamic. |
| // Tested with K=5 and operands with varying values. |
| func.func @batch_matmul_dynamic_reduction_size_B2_M1_N3_K5() { |
| %lhs = flow.tensor.dynamic_constant dense<[[[1.0, 2.0, 3.0, 4.0, 5.0]], |
| [[6.0, 7.0, 8.0, 9.0, 10.0]]]> |
| : tensor<2x1x5xf16> -> tensor<2x1x?xf16> |
| %rhs = flow.tensor.dynamic_constant dense<[[[11.0, 12.0, 13.0], |
| [14.0, 15.0, 16.0], |
| [17.0, 18.0, 19.0], |
| [20.0, 21.0, 22.0], |
| [23.0, 24.0, 25.0]], |
| [[26.0, 27.0, 28.0], |
| [29.0, 30.0, 31.0], |
| [32.0, 33.0, 34.0], |
| [35.0, 36.0, 37.0], |
| [38.0, 39.0, 40.0]]]> |
| : tensor<2x5x3xf16> -> tensor<2x?x3xf16> |
| |
| %cst = arith.constant 0.000000 : f16 |
| %2 = tensor.empty() : tensor<2x1x3xf16> |
| %3 = linalg.fill ins(%cst : f16) outs(%2 : tensor<2x1x3xf16>) -> tensor<2x1x3xf16> |
| %observed = linalg.generic { |
| indexing_maps = [#map, #map1, #map2], |
| iterator_types = ["parallel", "parallel", "parallel", "reduction"]} |
| ins(%lhs, %rhs : tensor<2x1x?xf16>, tensor<2x?x3xf16>) outs(%3 : tensor<2x1x3xf16>) { |
| ^bb0(%in: f16, %in_0: f16, %out: f16): |
| %6 = arith.mulf %in, %in_0 : f16 |
| %7 = arith.addf %out, %6 : f16 |
| linalg.yield %7 : f16 |
| } -> tensor<2x1x3xf16> |
| |
| %expected = flow.tensor.dynamic_constant dense<[ |
| [[285.0, 300.0, 315.0]], |
| [[1310.0, 1350.0, 1390.0]] |
| ]> : tensor<2x1x3xf16> -> tensor<2x1x3xf16> |
| check.expect_almost_eq(%observed, %expected, atol 1.0e-04) : tensor<2x1x3xf16> |
| return |
| } |
| |
| |
| // Softmax operation with dynamic reduction size. |
| // Number of samples: 2. Size of reduction dimension: 4. |
| func.func @softmax_dynamic_reduction_N2_K4(){ |
| |
| // Computed with numpy, values of ln([[3, 2, 4, 1.],[1, 7, 1, 1]]). |
| %input = flow.tensor.dynamic_constant dense< |
| [[1.09861229, 0.69314718, 1.38629436, 0. ], |
| [0. , 1.94591015, 0. , 0. ]]> : tensor<2x4xf32> -> tensor<2x?xf32> |
| |
| %expected = flow.tensor.dynamic_constant dense< |
| [[0.3, 0.2, 0.4, 0.1], |
| [0.1, 0.7, 0.1, 0.1]]> : tensor<2x4xf32> -> tensor<2x?xf32> |
| |
| %c_1_index = arith.constant 1 : index |
| %dim_0 = tensor.dim %input, %c_1_index : tensor<2x?xf32> |
| %output = tensor.empty(%dim_0) : tensor<2x?xf32> |
| |
| %sm = linalg.softmax dimension(1) ins(%input : tensor<2x?xf32>) |
| outs(%output : tensor<2x?xf32>) -> tensor<2x?xf32> |
| |
| check.expect_almost_eq(%sm, %expected, atol 1.0e-04) : tensor<2x?xf32> |
| return |
| } |
| |
| // Softmax operation with dynamic reduction size. |
| // Number of samples: 1. Size of reduction dimension: 1. |
| func.func @softmax_dynamic_reduction_N1_K1(){ |
| |
| %input = flow.tensor.dynamic_constant dense<[[-77.7]]> : tensor<1x1xf32> -> tensor<1x?xf32> |
| %expected = flow.tensor.dynamic_constant dense<[[1.0]]> : tensor<1x1xf32> -> tensor<1x?xf32> |
| %c_1_index = arith.constant 1 : index |
| %dim_0 = tensor.dim %input, %c_1_index : tensor<1x?xf32> |
| %output = tensor.empty(%dim_0) : tensor<1x?xf32> |
| %sm = linalg.softmax dimension(1) ins(%input : tensor<1x?xf32>) |
| outs(%output : tensor<1x?xf32>) -> tensor<1x?xf32> |
| |
| check.expect_almost_eq(%sm, %expected, atol 1.0e-04) : tensor<1x?xf32> |
| return |
| } |
| |
| |
| // Softmax operation with dynamic reduction size. |
| // Number of samples: 65. Size of reduction dimension: 1531. |
| func.func @softmax_dynamic_reduction_N65_K1531(){ |
| |
| %input = flow.tensor.dynamic_constant dense<-3.1415> : |
| tensor<65x1531xf32> -> tensor<65x?xf32> |
| |
| // 1/1531 is 0.0006531678641410843 |
| %expected = flow.tensor.dynamic_constant dense<0.0006531678641410843> : |
| tensor<65x1531xf32> -> tensor<65x?xf32> |
| |
| %c_1_index = arith.constant 1 : index |
| %dim_0 = tensor.dim %input, %c_1_index : tensor<65x?xf32> |
| %output = tensor.empty(%dim_0) : tensor<65x?xf32> |
| |
| %sm = linalg.softmax dimension(1) ins(%input : tensor<65x?xf32>) |
| outs(%output : tensor<65x?xf32>) -> tensor<65x?xf32> |
| |
| check.expect_almost_eq(%sm, %expected, atol 1.0e-04) : tensor<65x?xf32> |
| return |
| } |
| |
| // Batch matmul with: |
| // |
| // B = 1 <- dynamic, chosen at runtime |
| // M = 1 <- dynamic, chosen at runtime |
| // N = 32 |
| // K = 64 |
| // |
| // Followed by an erf non-linearity. |
| // This previously failed with compilation error, see issue |
| // https://github.com/iree-org/iree/issues/21813 |
| func.func @batch_matmul_B1_M1_N32_K64_followed_by_erf() { |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c0_f32 = arith.constant 0.0 : f32 |
| |
| %lhs = flow.tensor.dynamic_constant dense<0.125> : tensor<1x1x64xf32> -> tensor<?x?x64xf32> |
| %rhs = flow.tensor.dynamic_constant dense<0.125> : tensor<1x64x32xf32> -> tensor<?x64x32xf32> |
| %B = tensor.dim %lhs, %c0 : tensor<?x?x64xf32> |
| %M = tensor.dim %lhs, %c1 : tensor<?x?x64xf32> |
| |
| %empty = tensor.empty(%B, %M) : tensor<?x?x32xf32> |
| %filled = linalg.fill ins(%c0_f32: f32) outs(%empty : tensor<?x?x32xf32>) -> tensor<?x?x32xf32> |
| %bmm = linalg.batch_matmul ins(%lhs, %rhs : tensor<?x?x64xf32>, tensor<?x64x32xf32>) |
| outs(%filled : tensor<?x?x32xf32>) -> tensor<?x?x32xf32> |
| |
| %erfed = math.erf %bmm : tensor<?x?x32xf32> |
| |
| // The reduction size K=64, and the values in the LHS and RHS tensors are constants |
| // 1/8 and 1/8. So the result of the batched matmul is 32 * (1/4) * (1/8) = 1.0. |
| // |
| // erf(1.0) = 0.8427007929497149 (check with python `math.erf(1.0)`) |
| %expected = flow.tensor.dynamic_constant dense<0.84270079> : tensor<1x1x32xf32> -> tensor<?x?x32xf32> |
| check.expect_almost_eq(%erfed, %expected, atol 1.0e-02) : tensor<?x?x32xf32> |
| return |
| } |
