llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/test/Dialect/TMTensor/convert_to_loops.mlir

369 lines
17 KiB
MLIR
Raw Normal View History

Re-organize project structure to separate PyTorch dependencies from core project. (#2542) This is a first step towards the structure we discussed here: https://gist.github.com/stellaraccident/931b068aaf7fa56f34069426740ebf20 There are two primary goals: 1. Separate the core project (C++ dialects and conversions) from the hard PyTorch dependencies. We move all such things into projects/pt1 as a starting point since they are presently entangled with PT1-era APIs. Additional work can be done to disentangle components from that (specifically LTC is identified as likely ultimately living in a `projects/ltc`). 2. Create space for native PyTorch2 Dynamo-based infra to be upstreamed without needing to co-exist with the original TorchScript path. Very little changes in this path with respect to build layering or options. These can be updated in a followup without commingling directory structure changes. This also takes steps toward a couple of other layering enhancements: * Removes the llvm-external-projects/torch-mlir-dialects sub-project, collapsing it into the main tree. * Audits and fixes up the core C++ build to account for issues found while moving things. This is just an opportunistic pass through but roughly ~halves the number of build actions for the project from the high 4000's to the low 2000's. It deviates from the discussed plan by having a `projects/` tree instead of `compat/`. As I was thinking about it, this will better accommodate the follow-on code movement. Once things are roughly in place and the CI passing, followups will focus on more in-situ fixes and cleanups.
2023-11-03 10:45:55 +08:00
// RUN: torch-mlir-opt -split-input-file -tm-tensor-to-loops %s | FileCheck %s
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scan_1d_inclusive(%0: memref<128xi32>, %1: memref<128xi32>) {
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%c0 = memref.alloc() : memref<i32>
tm_tensor.scan dimension(0) inclusive(true)
ins(%0 : memref<128xi32>) outs(%1, %c0 : memref<128xi32>, memref<i32>) {
^bb0(%arg0 : i32, %arg1 : i32):
%sum = arith.addi %arg0, %arg1 : i32
tm_tensor.yield %sum : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scan_1d_inclusive
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[BUFI:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[BUFO:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[ACC:.+]] = memref.alloc() : memref<i32>
// CHECK: scf.for %[[ARG1:.+]] = %[[C0]] to %[[C128]] step %[[C1]]
// CHECK: %[[COND:.+]] = arith.cmpi eq, %[[ARG1]], %[[C0]] : index
// CHECK: scf.if %[[COND]] {
// CHECK: %[[V1:.+]] = memref.load %[[BUFI]][%[[ARG1]]]
// CHECK: memref.store %[[V1]], %[[BUFO]][%[[ARG1]]]
// CHECK: } else {
// CHECK: %[[T1:.+]] = arith.subi %[[ARG1]], %[[C1]] : index
// CHECK: %[[V2:.+]] = memref.load %[[BUFO]][%[[T1]]]
// CHECK: %[[V3:.+]] = memref.load %[[BUFI]][%[[ARG1]]]
// CHECK: %[[V4:.+]] = arith.addi %[[V2]], %[[V3]] : i32
// CHECK: memref.store %[[V4]], %[[BUFO]][%[[ARG1]]]
// CHECK: memref.store %[[V4]], %[[ACC]][]
// CHECK: }
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scan_1d_exclusive(%0: memref<128xi32>, %1: memref<128xi32>) {
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%c0 = memref.alloc() : memref<i32>
tm_tensor.scan dimension(0) inclusive(false)
ins(%0 : memref<128xi32>) outs(%1, %c0 : memref<128xi32>, memref<i32>) {
^bb0(%arg0 : i32, %arg1 : i32):
%sum = arith.addi %arg0, %arg1 : i32
tm_tensor.yield %sum : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scan_1d_exclusive
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[BUFI:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[BUFO:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[ACC:.+]] = memref.alloc() : memref<i32>
// CHECK: scf.for %[[ARG1:.+]] = %[[C0]] to %[[C128]] step %[[C1]]
// CHECK: %[[COND:.+]] = arith.cmpi eq, %[[ARG1]], %[[C0]] : index
// CHECK: scf.if %[[COND]] {
// CHECK: %[[V0:.+]] = memref.load %[[ACC]][] : memref<i32>
// CHECK: memref.store %[[V0]], %[[BUFO]][%[[ARG1]]]
// CHECK: } else {
// CHECK: %[[T1:.+]] = arith.subi %[[ARG1]], %[[C1]] : index
// CHECK: %[[V2:.+]] = memref.load %[[BUFO]][%[[T1]]]
// CHECK: %[[V3:.+]] = memref.load %[[BUFI]][%[[T1]]]
// CHECK: %[[V4:.+]] = arith.addi %[[V2]], %[[V3]] : i32
// CHECK: memref.store %[[V4]], %[[BUFO]][%[[ARG1]]]
// CHECK: memref.store %[[V4]], %[[ACC]][]
// CHECK: }
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scan_2d(%0: memref<16x32xi32>, %1: memref<16x32xi32>) {
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%t0 = memref.alloc() : memref<32xi32>
tm_tensor.scan dimension(0) inclusive(true)
ins(%0 : memref<16x32xi32>) outs(%1, %t0 : memref<16x32xi32>, memref<32xi32>) {
^bb0(%arg0 : i32, %arg1 : i32):
%sum = arith.addi %arg0, %arg1 : i32
tm_tensor.yield %sum : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scan_2d
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[BUFI:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[BUFO:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[ACC:.+]] = memref.alloc() : memref<32xi32>
// CHECK: scf.for %[[ARG1:.+]] = %[[C0]] to %[[C16]] step %[[C1]]
// CHECK: scf.for %[[ARG2:.+]] = %[[C0]] to %[[C32]] step %[[C1]]
// CHECK: %[[COND:.+]] = arith.cmpi eq, %[[ARG1]], %[[C0]] : index
// CHECK: scf.if %[[COND]] {
// CHECK: %[[V1:.+]] = memref.load %[[BUFI]][%[[ARG1]], %[[ARG2]]]
// CHECK: memref.store %[[V1]], %[[BUFO]][%[[ARG1]], %[[ARG2]]]
// CHECK: } else {
// CHECK: %[[T1:.+]] = arith.subi %[[ARG1]], %[[C1]] : index
// CHECK: %[[V2:.+]] = memref.load %[[BUFO]][%[[T1]], %[[ARG2]]]
// CHECK: %[[V3:.+]] = memref.load %[[BUFI]][%[[ARG1]], %[[ARG2]]]
// CHECK: %[[V4:.+]] = arith.addi %[[V2]], %[[V3]] : i32
// CHECK: memref.store %[[V4]], %[[BUFO]][%[[ARG1]], %[[ARG2]]]
// CHECK: memref.store %[[V4]], %[[ACC]][%[[ARG2]]]
// CHECK: }
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_update_scalar_1D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<8xi32>, %indices: memref<3x1xi32>,
%updates: memref<3xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<3xi32>, memref<3x1xi32>)
outs(%original : memref<8xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
tm_tensor.yield %arg0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scatter_update_scalar_1D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[C3]] step %[[C1]] {
// CHECK: %[[T1:.+]] = memref.load %[[UPDATES]][%[[I]]] : memref<3xi32>
// CHECK: %[[T2:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]] : memref<3x1xi32>
// CHECK: %[[IDX:.+]] = arith.index_cast %[[T2]] : i32 to index
// CHECK: memref.store %[[T1]], %[[ORIGINAL]][%[[IDX]]]
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_add_scalar_2D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<4x3xi32>, %indices: memref<3x2xi32>,
%updates: memref<3xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0, 1>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<3xi32>, memref<3x2xi32>)
outs(%original : memref<4x3xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
%0 = arith.addi %arg1, %arg0 : i32
tm_tensor.yield %0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scatter_add_scalar_2D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[C3]] step %[[C1]] {
// CHECK: %[[T1:.+]] = memref.load %[[UPDATES]][%[[I]]] : memref<3xi32>
// CHECK: %[[T2:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]] : memref<3x2xi32>
// CHECK: %[[IDX1:.+]] = arith.index_cast %[[T2]] : i32 to index
// CHECK: %[[T3:.+]] = memref.load %[[INDICES]][%[[I]], %[[C1]]] : memref<3x2xi32>
// CHECK: %[[IDX2:.+]] = arith.index_cast %[[T3]] : i32 to index
// CHECK: %[[ORI:.+]] = memref.load %[[ORIGINAL]][%[[IDX1]], %[[IDX2]]] : memref<4x3xi32>
// CHECK: %[[ADD:.+]] = arith.addi %[[ORI]], %[[T1]] : i32
// CHECK: memref.store %[[ADD]], %[[ORIGINAL]][%[[IDX1]], %[[IDX2]]]
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_update_slice_2D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<4x3xi32>, %indices: memref<2x1xi32>,
%updates: memref<2x3xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<2x3xi32>, memref<2x1xi32>)
outs(%original : memref<4x3xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
tm_tensor.yield %arg0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK: func.func @scatter_update_slice_2D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[C2]] step %[[C1]] {
// CHECK: scf.for %[[J:.+]] = %[[C0]] to %[[C3]] step %[[C1]] {
// CHECK: %[[UPDATE:.+]] = memref.load %[[UPDATES]][%[[I]], %[[J]]]
// CHECK: %[[INDEX:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]]
// CHECK: %[[LOC:.+]] = arith.index_cast %[[INDEX]] : i32 to index
// CHECK: memref.store %[[UPDATE]], %[[ORIGINAL]][%[[LOC]], %[[J]]]
// CHECK: }
// CHECK: }
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_add_scalar_1D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<8xi32>, %indices: memref<3x1xi32>,
%updates: memref<3xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<3xi32>, memref<3x1xi32>)
outs(%original : memref<8xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
%0 = arith.addi %arg1, %arg0 : i32
tm_tensor.yield %0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scatter_add_scalar_1D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[C3]] step %[[C1]] {
// CHECK: %[[T1:.+]] = memref.load %[[UPDATES]][%[[I]]] : memref<3xi32>
// CHECK: %[[T2:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]] : memref<3x1xi32>
// CHECK: %[[IDX:.+]] = arith.index_cast %[[T2]] : i32 to index
// CHECK: %[[ORI:.+]] = memref.load %[[ORIGINAL]][%[[IDX]]] : memref<8xi32>
// CHECK: %[[ADD:.+]] = arith.addi %[[ORI]], %[[T1]] : i32
// CHECK: memref.store %[[ADD]], %[[ORIGINAL]][%[[IDX]]]
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_add_slice_2D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<4x3xi32>, %indices: memref<2x1xi32>,
%updates: memref<2x3xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<2x3xi32>, memref<2x1xi32>)
outs(%original : memref<4x3xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
%0 = arith.addi %arg1, %arg0 : i32
tm_tensor.yield %0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK: func.func @scatter_add_slice_2D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[C2]] step %[[C1]] {
// CHECK: scf.for %[[J:.+]] = %[[C0]] to %[[C3]] step %[[C1]] {
// CHECK: %[[UPDATEVAL:.+]] = memref.load %[[UPDATES]][%[[I]], %[[J]]]
// CHECK: %[[INDEXVAL:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]]
// CHECK: %[[INDEX:.+]] = arith.index_cast %[[INDEXVAL]] : i32 to index
// CHECK: %[[ORIGINALVAL:.+]] = memref.load %[[ORIGINAL]][%[[INDEX]], %[[J]]]
// CHECK: %[[STOREVAL:.+]] = arith.addi %[[ORIGINALVAL]], %[[UPDATEVAL]]
// CHECK: memref.store %[[STOREVAL]], %[[ORIGINAL]][%[[INDEX]], %[[J]]]
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_update_scalar_dynamic_1D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<?xi32>, %indices: memref<?x1xi32>,
%updates: memref<?xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<?xi32>, memref<?x1xi32>)
outs(%original : memref<?xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
tm_tensor.yield %arg0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scatter_update_scalar_dynamic_1D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[UB:.+]] = memref.dim %[[UPDATES]], %[[C0]] : memref<?xi32>
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[UB]] step %[[C1]] {
// CHECK: %[[T1:.+]] = memref.load %[[UPDATES]][%[[I]]] : memref<?xi32>
// CHECK: %[[T2:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]] : memref<?x1xi32>
// CHECK: %[[IDX:.+]] = arith.index_cast %[[T2]] : i32 to index
// CHECK: memref.store %[[T1]], %[[ORIGINAL]][%[[IDX]]]
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_add_scalar_dynamic_2D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<?x?xi32>, %indices: memref<?x2xi32>,
%updates: memref<?xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0, 1>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<?xi32>, memref<?x2xi32>)
outs(%original : memref<?x?xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
%0 = arith.addi %arg1, %arg0 : i32
tm_tensor.yield %0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scatter_add_scalar_dynamic_2D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[UB:.+]] = memref.dim %[[UPDATES]], %[[C0]] : memref<?xi32>
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[UB]] step %[[C1]] {
// CHECK: %[[T1:.+]] = memref.load %[[UPDATES]][%[[I]]] : memref<?xi32>
// CHECK: %[[T2:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]] : memref<?x2xi32>
// CHECK: %[[IDX1:.+]] = arith.index_cast %[[T2]] : i32 to index
// CHECK: %[[T3:.+]] = memref.load %[[INDICES]][%[[I]], %[[C1]]] : memref<?x2xi32>
// CHECK: %[[IDX2:.+]] = arith.index_cast %[[T3]] : i32 to index
// CHECK: %[[ORI:.+]] = memref.load %[[ORIGINAL]][%[[IDX1]], %[[IDX2]]] : memref<?x?xi32>
// CHECK: %[[ADD:.+]] = arith.addi %[[ORI]], %[[T1]] : i32
// CHECK: memref.store %[[ADD]], %[[ORIGINAL]][%[[IDX1]], %[[IDX2]]]
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_update_slice_dynamic_2D(
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
%original: memref<?x?xi32>, %indices: memref<?x1xi32>,
%updates: memref<?x?xi32>) {
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
tm_tensor.scatter {dimension_map= array<i64: 0>} unique_indices(true)
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
ins(%updates, %indices : memref<?x?xi32>, memref<?x1xi32>)
outs(%original : memref<?x?xi32>) {
^bb0(%arg0: i32, %arg1: i32): // no predecessors
tm_tensor.yield %arg0 : i32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK: func.func @scatter_update_slice_dynamic_2D
Add TMTensor dialect to torch-mlir This is intended to explore support for non-structured ops that can't be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter` are added as the first such ops. The dialect should aim to be upstreamed in the future.
2022-02-03 07:01:38 +08:00
// CHECK-SAME: %[[ORIGINAL:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INDICES:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[UPDATES:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[UB1:.+]] = memref.dim %[[UPDATES]], %[[C0]] : memref<?x?xi32>
// CHECK-DAG: %[[UB2:.+]] = memref.dim %[[UPDATES]], %[[C1]] : memref<?x?xi32>
// CHECK: scf.for %[[I:.+]] = %[[C0]] to %[[UB1]] step %[[C1]] {
// CHECK: scf.for %[[J:.+]] = %[[C0]] to %[[UB2]] step %[[C1]] {
// CHECK: %[[UPDATEVAL:.+]] = memref.load %[[UPDATES]][%[[I]], %[[J]]]
// CHECK: %[[INDEXVAL:.+]] = memref.load %[[INDICES]][%[[I]], %[[C0]]]
// CHECK: %[[INDEX:.+]] = arith.index_cast %[[INDEXVAL]] : i32 to index
// CHECK: memref.store %[[UPDATEVAL]], %[[ORIGINAL]][%[[INDEX]], %[[J]]]
Extend tm_tensor.scatter op semantic to carry unique_indices attribute There are cases where the op may update the same indices multiple times. In this context, we can not parallelize updates. Instead, we have to execute them sequentially. Adding a boolean attribute to control the behavior. Also adding test cases for invalid IR.
2022-02-21 15:40:02 +08:00
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @scatter_partial_slices(%arg0: memref<2x64x12xf32>, %arg1: memref<2x3xi32>, %arg2: memref<2x1x12xf32>) {
Extend tm_tensor.scatter op semantic to carry unique_indices attribute There are cases where the op may update the same indices multiple times. In this context, we can not parallelize updates. Instead, we have to execute them sequentially. Adding a boolean attribute to control the behavior. Also adding test cases for invalid IR.
2022-02-21 15:40:02 +08:00
tm_tensor.scatter
[torch] Rework lowering to tm_tensor.scatter to stop serialization (#2940) We collapsed and broadcasted scatter indices to a single element version. We should instead upport `tm_tensor.scatter`s support for multiple indices and the implicitly broadcasted behavior. This avoids the serialization and materializing a needlessly large indices tensor.
2024-02-28 03:46:57 +08:00
{dimension_map= array<i64: 0, 1, 2>}
Extend tm_tensor.scatter op semantic to carry unique_indices attribute There are cases where the op may update the same indices multiple times. In this context, we can not parallelize updates. Instead, we have to execute them sequentially. Adding a boolean attribute to control the behavior. Also adding test cases for invalid IR.
2022-02-21 15:40:02 +08:00
unique_indices(true)
ins(%arg2, %arg1 : memref<2x1x12xf32>, memref<2x3xi32>)
outs(%arg0 : memref<2x64x12xf32>) {
^bb0(%arg3: f32, %arg4: f32):
tm_tensor.yield %arg4 : f32
}
return
}
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @scatter_partial_slices
Extend tm_tensor.scatter op semantic to carry unique_indices attribute There are cases where the op may update the same indices multiple times. In this context, we can not parallelize updates. Instead, we have to execute them sequentially. Adding a boolean attribute to control the behavior. Also adding test cases for invalid IR.
2022-02-21 15:40:02 +08:00
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]
Fix convert_to_loops.mlir format
2022-03-16 06:05:12 +08:00
// CHECK-DAG: %[[C0:.+]] = arith.constant
// CHECK-DAG: %[[C1:.+]] = arith.constant
// CHECK-DAG: %[[C2:.+]] = arith.constant
// CHECK-DAG: %[[C12:.+]] = arith.constant
// CHECK: scf.for %[[ARG3:.+]] = %[[C0]] to %[[C2]] step %[[C1]] {
// CHECK-NEXT: scf.for %[[ARG4:.+]] = %[[C0]] to %[[C1]] step %[[C1]] {
// CHECK-NEXT: scf.for %[[ARG5:.+]] = %[[C0]] to %[[C12]] step %[[C1]] {
// CHECK-NEXT: %[[LOAD0:.+]] = memref.load %[[ARG1]][%[[ARG3]], %[[C0]]] : memref<2x3xi32>
// CHECK-NEXT: %[[CAST0:.+]] = arith.index_cast %[[LOAD0]] : i32 to index
// CHECK-NEXT: %[[LOAD1:.+]] = memref.load %[[ARG1]][%[[ARG3]], %[[C1]]] : memref<2x3xi32>
// CHECK-NEXT: %[[CAST1:.+]] = arith.index_cast %[[LOAD1]] : i32 to index
// CHECK-NEXT: %[[ADD1:.+]] = arith.addi %[[CAST1]], %[[ARG4]] : index
// CHECK-NEXT: %[[LOAD2:.+]] = memref.load %[[ARG1]][%[[ARG3]], %[[C2]]] : memref<2x3xi32>
// CHECK-NEXT: %[[CAST2:.+]] = arith.index_cast %[[LOAD2]] : i32 to index
// CHECK-NEXT: %[[ADD2:.+]] = arith.addi %[[CAST2]], %[[ARG5]] : index
// CHECK-NEXT: %[[LOAD3:.+]] = memref.load %[[ARG0]][%[[CAST0]], %[[ADD1]], %[[ADD2]]] : memref<2x64x12xf32>
// CHECK-NEXT: memref.store %[[LOAD3]], %[[ARG0]][%[[CAST0]], %[[ADD1]], %[[ADD2]]] : memref<2x64x12xf32>