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/TCP/ops.mlir

23 lines
882 B
MLIR
Raw Normal View History

[RefBackend] Split out RefBackend (refback) dialect from TCP. This is the first in a patch series that is refactoring the constellation of things variously called or associated with "E2E", "RefE2E", "npcomprt", and "TCP" into a more cleanly layered result. Concretely, this first patch fixes the fact that TCP was basically acting like a dumping ground needed by the reference backend. This splits it out, which is fairly mechanical, but touches a lot of lines of code (basically replacing `tcp` with `refback` and `TCP` with `RefBackend). Now, the RefBackend dialect is that dumping ground, which is slighly better, as it starts allowing TCP to become a nice clean middle layer that is not related per se to the reference backend. The previous name RefE2E or "reference e2e flow" was super confusing. Now that we are seeing more clearly where the "backend" distinction lies, the [RefBackend] commit tag is born :)
2020-10-07 06:44:18 +08:00
// RUN: npcomp-opt <%s | npcomp-opt | FileCheck %s
[RefE2E] Add support for "max". This cleans up the lowering pipeline to easily allow extending to multiple binary ops. It looks fairly repetitive at multiple levels, but I don't want to prematurely generalize. I think that in principle we could derive a large swatch of TCF + TCP from a single linalg-style specification. Another direction is to use an OpInterface (something like "buildLinalgGenericBody"). I'm keeping my eye on it. In a subsequent commit, I'll mechanically add a set of binary ops modeled off of the std arithmetic ops.
2020-09-22 10:14:27 +08:00
[TCP] Replace tcp.matmul with linalg.matmul. This involved adding a `tcp.splatted` op to splat a dynamically sized init tensor. See rationale in TCPOps.td docs. One interesting observation is that when lowering tcf.matmul to linalg.matmul, we need to both 1) create the error checks and 2) calculate a shape transfer function to create the init tensors. Previously, 2) was deferred to bufferizing tcp.matmul later. I'm not sure if this is a conflation of concerns or not. For now, it's not a big burden.
2020-11-10 07:49:22 +08:00
// CHECK-LABEL: @broadcast_to
func @broadcast_to(%arg0: tensor<?xf32>, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
// CHECK: tcp.broadcast_to
%0 = tcp.broadcast_to %arg0, %arg1 : (tensor<?xf32>, tensor<?xindex>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32>
}
// CHECK-LABEL: @splatted
func @splatted(%arg0: f32, %arg1: tensor<?xindex>) -> tensor<?x?xf32> {
// CHECK: tcp.splatted
%0 = tcp.splatted %arg0, %arg1 : (f32, tensor<?xindex>) -> tensor<?x?xf32>
[RefE2E] Add support for matmul. I'm pretty happy with how this turned out. It looks pretty much like it should -- one change at each layer. This particular op bottoms out on linalg which takes care of the rest. - Add tcf.matmul - Add tcp.matmul - Add TCF->TCP lowering - Add tcp.matmul shape transfer function (BypassShapes.cpp) - Add tcp.matmul -> linalg.matmul lowering (LowerShapedResultsToMemref.cpp) - Add support to LowerShapeConstraints for lowering the new shape.cstr_require This matmul op is pretty limited in its capabilities. There is no batching and no multidimensional contraction. Certainly more design work will be needed to find the right abstractions that aren't too general but also help to canonicalize many cases from frontends. This is mainly to show that adding a new op needn't be very "scary" once we have the e2e infra in place. Also, - this clears out some exploratory cruft from the TCF dialect now that this is starting to become real.
2020-09-18 09:56:01 +08:00
return %0 : tensor<?x?xf32>
}
Import TCP pad
2020-12-18 02:56:46 +08:00
// CHECK-LABEL: @pad
func @pad(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?xindex>, %arg2: tensor<?xindex>, %arg3: f32) -> tensor<?x?x?x?xf32> {
// CHECK: tcp.pad
%0 = tcp.pad %arg0, %arg1, %arg2, %arg3 : (tensor<?x?x?x?xf32>, tensor<?xindex>, tensor<?xindex>, f32) -> tensor<?x?x?x?xf32>
return %0 : tensor<?x?x?x?xf32>
}