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

10 lines
474 B
MLIR
Raw Normal View History

Add a torch.kernel_call op and associated predicates.
2020-09-30 05:17:34 +08:00
// RUN: npcomp-opt %s | npcomp-opt | FileCheck %s
func @kernel_call(%arg0 : si32, %arg1 : tensor<3x4xf32>) -> tensor<*xf32> {
// CHECK: %0 = torch.kernel_call "somens::someop" %arg0, %arg1 : (si32, tensor<3x4xf32>) -> tensor<*xf32>
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
%1 = torch.kernel_call "somens::someop" %arg0, %arg1 : (si32, tensor<3x4xf32>) -> (tensor<*xf32>) {
sigArgTypes = [], sigRetTypes = [], sigIsVararg = false, sigIsVarret = false, sigIsMutable = false
}
Add a torch.kernel_call op and associated predicates.
2020-09-30 05:17:34 +08:00
return %1 : tensor<*xf32>
}