| // RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(torch-iree-tm-tensor-to-linalg-ext))" %s | FileCheck %s |
| |
| func.func @scatter_update( |
| %original: tensor<8xi32>, %indices: tensor<3x1xi32>, |
| %updates: tensor<3xi32>) -> tensor<8xi32> { |
| %0 = tm_tensor.scatter {dimension_map = array<i64: 0>} unique_indices(true) |
| ins(%updates, %indices : tensor<3xi32>, tensor<3x1xi32>) |
| outs(%original : tensor<8xi32>) { |
| ^bb0(%update: i32, %orig: i32): // no predecessors |
| tm_tensor.yield %update: i32 |
| } -> tensor<8xi32> |
| return %0 : tensor<8xi32> |
| } |
| // CHECK-LABEL: func.func @scatter_update( |
| // CHECK-SAME: %[[ORIGINAL:.*]]: tensor<8xi32>, |
| // CHECK-SAME: %[[INDICES:.*]]: tensor<3x1xi32>, |
| // CHECK-SAME: %[[UPDATES:.*]]: tensor<3xi32>) -> tensor<8xi32> { |
| // CHECK: %[[SCATTER:.*]] = iree_linalg_ext.scatter |
| // CHECK-SAME: dimension_map = [0] |
| // CHECK-SAME: unique_indices(true) |
| // CHECK-SAME: ins(%[[UPDATES]], %[[INDICES]] : tensor<3xi32>, tensor<3x1xi32>) |
| // CHECK-SAME: outs(%[[ORIGINAL]] : tensor<8xi32>) { |
| // CHECK: ^bb0(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32): |
| // CHECK: iree_linalg_ext.yield %[[ARG0]] : i32 |
| // CHECK: } -> tensor<8xi32> |
| // CHECK: return %[[SCATTER:.*]] : tensor<8xi32> |
| |
| // ----- |
| func.func @scatter_add( |
| %original: tensor<8xi32>, %indices: tensor<3x1xi32>, |
| %updates: tensor<3xi32>) -> tensor<8xi32> { |
| %0 = tm_tensor.scatter {dimension_map = array<i64: 0>} unique_indices(true) |
| ins(%updates, %indices : tensor<3xi32>, tensor<3x1xi32>) |
| outs(%original : tensor<8xi32>) { |
| ^bb0(%update: i32, %orig: i32): // no predecessors |
| %add = arith.addi %orig, %update: i32 |
| tm_tensor.yield %add: i32 |
| } -> tensor<8xi32> |
| return %0 : tensor<8xi32> |
| } |
| // CHECK-LABEL: func.func @scatter_add( |
| // CHECK-SAME: %[[ORIGINAL:.*]]: tensor<8xi32>, |
| // CHECK-SAME: %[[INDICES:.*]]: tensor<3x1xi32>, |
| // CHECK-SAME: %[[UPDATES:.*]]: tensor<3xi32>) -> tensor<8xi32> { |
| // CHECK: %[[SCATTER:.*]] = iree_linalg_ext.scatter |
| // CHECK-SAME: unique_indices(true) |
| // CHECK-SAME: ins(%[[UPDATES]], %[[INDICES]] : tensor<3xi32>, tensor<3x1xi32>) |
| // CHECK-SAME: outs(%[[ORIGINAL]] : tensor<8xi32>) { |
| // CHECK: ^bb0(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32): |
| // CHECK: %[[ADD:.*]] = arith.addi %[[ARG1]], %[[ARG0]] : i32 |
| // CHECK: iree_linalg_ext.yield %[[ADD]] : i32 |
| // CHECK: } -> tensor<8xi32> |
| // CHECK: return %[[SCATTER:.*]] : tensor<8xi32> |
| |
| // ----- |
