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-ranked-shapes.mlir

23 lines
756 B
MLIR
Raw Normal View History

Lower !shape.shape to SSA values. This uses an approach inspired by what is done in IREE. See comments on LowerRankedShapes.cpp for how it works. The basic gist is that we have an op that creates a !shape.shape from a set of SSA values representing the extents, and then iteratively replace any op producing a !shape.shape with instances of that op.
2020-05-12 12:22:40 +08:00
// RUN: npcomp-opt -lower-ranked-shapes <%s | FileCheck %s --dump-input=fail
// CHECK-LABEL: func @broadcast_rank2_rank1
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
func @broadcast_rank2_rank1(%arg0: index, %arg1: index, %arg2: index) -> (index, index) {
Lower !shape.shape to SSA values. This uses an approach inspired by what is done in IREE. See comments on LowerRankedShapes.cpp for how it works. The basic gist is that we have an op that creates a !shape.shape from a set of SSA values representing the extents, and then iteratively replace any op producing a !shape.shape with instances of that op.
2020-05-12 12:22:40 +08:00
// CHECK-NOT: shape.broadcast
// CHECK-NOT: tcp.get_extent
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
%0 = shape.from_extents %arg0, %arg1
%1 = shape.from_extents %arg2
Lower !shape.shape to SSA values. This uses an approach inspired by what is done in IREE. See comments on LowerRankedShapes.cpp for how it works. The basic gist is that we have an op that creates a !shape.shape from a set of SSA values representing the extents, and then iteratively replace any op producing a !shape.shape with instances of that op.
2020-05-12 12:22:40 +08:00
%2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
%e0 = tcp.get_extent %2, 0
%e1 = tcp.get_extent %2, 1
return %e0, %e1 : index, index
}
Rename tcp.abort_if to tcp.shape_observe_error This more clearly captures its semantics as a structural "observer" of code that we currently mark as NoSideEffect but eventually lowers to eager error handling code. Also, update LowerRankedShapes to erase it, now that the layering here is clear. That pass reifies the eager error handling code, so the need for the dummy op to keep things alive isn't needed. With this change, we are now ready to start lowering to LLVM! This is the current print-ir-after-all from e2e-lowering-pipeline: https://reviews.llvm.org/P8221
2020-05-19 04:35:25 +08:00
// CHECK-LABEL: func @erase_shape_observe_error
func @erase_shape_observe_error(%arg0: index) {
// CHECK-NOT tcp.shape_observe_error
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
%0 = shape.from_extents %arg0
Rename tcp.abort_if to tcp.shape_observe_error This more clearly captures its semantics as a structural "observer" of code that we currently mark as NoSideEffect but eventually lowers to eager error handling code. Also, update LowerRankedShapes to erase it, now that the layering here is clear. That pass reifies the eager error handling code, so the need for the dummy op to keep things alive isn't needed. With this change, we are now ready to start lowering to LLVM! This is the current print-ir-after-all from e2e-lowering-pipeline: https://reviews.llvm.org/P8221
2020-05-19 04:35:25 +08:00
"tcp.shape_observe_error"(%0) : (!shape.shape) -> none
return
}