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/E2E/lower-to-memref-abi.mlir

31 lines
1.2 KiB
MLIR
Raw Normal View History

Lower to the upstream memref ABI. Specifically, we use unranked memrefs which get passed as a fixed-size set of arguments/returns. One big caveat about this is that returning results isn't going to work. See TODO in LowerTensorLoadOp. This is far from enough runtime-wise, but it starts to demarcate a plausible layering. Notice for example how this removes the runtime-dependence from LowerRankedShapes. Eventually, we want to have an `npcomp_rt` or `npcomp_hal` dialect with its own set of runtime types that will supercede this. See comments in LowerTensorLoadOp for more direction about where this is going to evolve.
2020-05-16 07:33:01 +08:00
// RUN: npcomp-opt -lower-to-memref-abi <%s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @identity
func @identity(%arg0: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: return %arg0 : memref<*xf32>
return %arg0 : tensor<?xf32>
}
// CHECK-LABEL: func @basic(
// CHECK-SAME: %[[VAL_1:.*]]: memref<*xf32>) -> memref<*xf32> {
func @basic(%arg0: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: %[[VAL_2:.*]] = memref_cast %[[VAL_1]] : memref<*xf32> to memref<?xf32>
// CHECK: %[[VAL_3:.*]] = dim %[[VAL_2]], 0 : memref<?xf32>
Use upstream shape.from_extents Replace our local `tcp.shape_from_extents` op with the upstream `shape.from_extents` op.
2020-05-22 05:51:01 +08:00
// CHECK: %[[VAL_4:.*]] = shape.from_extents %[[VAL_3]]
Lower to the upstream memref ABI. Specifically, we use unranked memrefs which get passed as a fixed-size set of arguments/returns. One big caveat about this is that returning results isn't going to work. See TODO in LowerTensorLoadOp. This is far from enough runtime-wise, but it starts to demarcate a plausible layering. Notice for example how this removes the runtime-dependence from LowerRankedShapes. Eventually, we want to have an `npcomp_rt` or `npcomp_hal` dialect with its own set of runtime types that will supercede this. See comments in LowerTensorLoadOp for more direction about where this is going to evolve.
2020-05-16 07:33:01 +08:00
%shape = shape.shape_of %arg0 : tensor<?xf32>
// CHECK: %[[VAL_5:.*]] = tcp.alloc_memref %[[VAL_4]] : memref<?xf32>
%memref = tcp.alloc_memref %shape : memref<?xf32>
// CHECK: %[[VAL_6:.*]] = memref_cast %[[VAL_1]] : memref<*xf32> to memref<?xf32>
// CHECK: linalg.copy(%[[VAL_6]], %[[VAL_5]]) : memref<?xf32>, memref<?xf32>
tensor_store %arg0, %memref : memref<?xf32>
// CHECK: %[[VAL_7:.*]] = memref_cast %[[VAL_5]] : memref<?xf32> to memref<*xf32>
%ret = tensor_load %memref : memref<?xf32>
// CHECK: return %[[VAL_7]] : memref<*xf32>
return %ret: tensor<?xf32>
}