| // Regression test for dynamic batch_matmul with WMMAR3 on RDNA3 (gfx1100). | |
| // | |
| // WMMAR3 has accumulator layout outer={8,1} which requires an expand_shape | |
| // on the output. With dynamic shapes, tensor.dim users on the forall result | |
| // previously blocked the ExpandDestinationForallOp pattern, causing a | |
| // separate shared memory allocation for the output accumulator that exceeded | |
| // the 65536-byte limit. | |
| func.func @dynamic_batch_matmul_transposed_rhs() { | |
| %lhs = flow.tensor.dynamic_constant dense<1.0> : tensor<2x128x128xf16> -> tensor<2x?x128xf16> | |
| %rhs = flow.tensor.dynamic_constant dense<1.0> : tensor<2x128x128xf16> -> tensor<2x?x128xf16> | |
| %cst = arith.constant 0.0 : f32 | |
| %c1 = arith.constant 1 : index | |
| %m = tensor.dim %lhs, %c1 : tensor<2x?x128xf16> | |
| %n = tensor.dim %rhs, %c1 : tensor<2x?x128xf16> | |
| %init = tensor.empty(%m, %n) : tensor<2x?x?xf32> | |
| %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<2x?x?xf32>) -> tensor<2x?x?xf32> | |
| %observed = 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(%lhs, %rhs : tensor<2x?x128xf16>, tensor<2x?x128xf16>) | |
| outs(%fill : tensor<2x?x?xf32>) -> tensor<2x?x?xf32> | |
| // Each output element = sum(1.0 * 1.0, K=128) = 128.0 | |
| %expected = flow.tensor.dynamic_constant dense<128.0> : tensor<2x128x128xf32> -> tensor<2x?x?xf32> | |
| check.expect_almost_eq(%observed, %expected, atol 1.0e-01) : tensor<2x?x?xf32> | |
| return | |
| } |