| func.func @gather_from_splat_tensor() { |
| %source = util.unfoldable_constant dense<0> : tensor<10x10xi32> |
| %empty = tensor.empty() : tensor<1x10xi32> |
| %indices = util.unfoldable_constant dense<0> : tensor<1xi32> |
| %result = iree_linalg_ext.gather dimension_map = [0] |
| ins(%source, %indices : tensor<10x10xi32>, tensor<1xi32>) |
| outs(%empty : tensor<1x10xi32>) -> tensor<1x10xi32> |
| |
| check.expect_eq_const(%result, dense<0> : tensor<1x10xi32>) |
| : tensor<1x10xi32> |
| return |
| } |
| |
| func.func @gather_2d_index_with_batch() { |
| %source = util.unfoldable_constant dense<[[0, 1], [2, 3]]> : tensor<2x2xi32> |
| %empty = tensor.empty() : tensor<2xi32> |
| %indices = util.unfoldable_constant dense<[[0, 1], [1, 0]]> : tensor<2x2xi32> |
| %result = iree_linalg_ext.gather dimension_map = [0, 1] |
| ins(%source, %indices : tensor<2x2xi32>, tensor<2x2xi32>) |
| outs(%empty: tensor<2xi32>) -> tensor<2xi32> |
| check.expect_eq_const(%result, dense<[1, 2]> : tensor<2xi32>) : tensor<2xi32> |
| return |
| } |
| |
| func.func @gather_2d_index_no_batch() { |
| %source = util.unfoldable_constant dense<[[[0], [1]], [[0], [0]]]> : tensor<2x2x1xi32> |
| %empty = tensor.empty() : tensor<1xi32> |
| %indices = util.unfoldable_constant dense<[0, 1]> : tensor<2xi32> |
| %result = iree_linalg_ext.gather dimension_map = [0, 1] |
| ins(%source, %indices : tensor<2x2x1xi32>, tensor<2xi32>) |
| outs(%empty: tensor<1xi32>) -> tensor<1xi32> |
| check.expect_eq_const(%result, dense<[1]> : tensor<1xi32>) : tensor<1xi32> |
| return |
| } |
| |
| func.func @gather_1d_index_no_batch() { |
| %source = util.unfoldable_constant dense<[[0, 1], [2, 3]]> : tensor<2x2xi32> |
| %empty = tensor.empty() : tensor<2xi32> |
| %indices = util.unfoldable_constant dense<[1]> : tensor<1xi32> |
| %result = iree_linalg_ext.gather dimension_map = [0] |
| ins(%source, %indices : tensor<2x2xi32>, tensor<1xi32>) |
| outs(%empty: tensor<2xi32>) -> tensor<2xi32> |
| check.expect_eq_const(%result, dense<[2, 3]> : tensor<2xi32>) : tensor<2xi32> |
| return |
| } |
| |
| func.func @gather_perm_map() { |
| %source = util.unfoldable_constant dense<[[[0], [1]], [[2], [3]]]> : tensor<2x2x1xi32> |
| %empty = tensor.empty() : tensor<1xi32> |
| %indices = util.unfoldable_constant dense<[0, 1]> : tensor<2xi32> |
| %result = iree_linalg_ext.gather dimension_map = [1, 0] |
| ins(%source, %indices : tensor<2x2x1xi32>, tensor<2xi32>) |
| outs(%empty: tensor<1xi32>) -> tensor<1xi32> |
| check.expect_eq_const(%result, dense<[2]> : tensor<1xi32>) : tensor<1xi32> |
| return |
| } |
| |
| func.func @gather_fuse_elementwise() { |
| %cst = arith.constant 2 : i32 |
| %source = util.unfoldable_constant dense<[[0, 1], [2, 3]]> : tensor<2x2xi32> |
| %empty = tensor.empty() : tensor<2xi32> |
| %indices = util.unfoldable_constant dense<[1]> : tensor<1xi32> |
| %result = iree_linalg_ext.gather dimension_map = [0] |
| ins(%source, %indices : tensor<2x2xi32>, tensor<1xi32>) |
| outs(%empty: tensor<2xi32>) -> tensor<2xi32> |
| %generic = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%result : tensor<2xi32>) outs(%empty: tensor<2xi32>) { |
| ^bb0(%arg0: i32, %arg1: i32): |
| %0 = arith.muli %arg0, %cst : i32 |
| linalg.yield %0 : i32 |
| } -> tensor<2xi32> |
| check.expect_eq_const(%generic, dense<[4, 6]> : tensor<2xi32>) : tensor<2xi32> |
| return |
| } |
