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/adjust-calling-conventions....

47 lines
2.6 KiB
MLIR
Raw Normal View History

Add torch-adjust-calling-conventions pass. This pass incorporates torch.type_bound info and also removes NoneType returns (eventually it will rewrite tuple types too, but can't yet because !basicpy.TupleType doesn't track element types). Recommend looking at adjust-calling-conventions.mlir first to see what it is doing, and holding your nose for the implementation of the pass. I decided to implement this with the conversion framework, because it gives us *some* goodies for type conversion -- mainly avoiding large amounts of tricky RAUW dances. Unfortunately, the conversion framework isn't a perfect fit for a couple reasons: - the incorporation of torch.type_bound is a context-sensitive rewrite (requires looking at the arg attr, not just the type). - NoneType conversion is 1->0, which requires some special handling - (not implemented yet) 1->N tuple type conversions require special handling. It's a little bit scary, but on balance doing it the other way would have its own downsides.
2021-04-02 08:36:18 +08:00
// RUN: npcomp-opt -torch-adjust-calling-conventions -allow-unregistered-dialect -split-input-file %s | FileCheck %s
// CHECK-LABEL: func @basic(
// CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<[2,3,?]:f32>) -> !numpy.ndarray<*:!numpy.any_dtype> {
// CHECK: %[[RET:.*]] = numpy.static_info_cast %[[ARG]] : !numpy.ndarray<[2,3,?]:f32> to !numpy.ndarray<*:!numpy.any_dtype>
// CHECK: return %[[RET]] : !numpy.ndarray<*:!numpy.any_dtype>
func @basic(%arg0: !numpy.ndarray<*:!numpy.any_dtype> {torch.type_bound = !numpy.ndarray<[2,3,?]:f32>}) -> !numpy.ndarray<*:!numpy.any_dtype> {
return %arg0 : !numpy.ndarray<*:!numpy.any_dtype>
}
// CHECK-LABEL: func @no_type_bound(
// CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<*:!numpy.any_dtype>) -> !numpy.ndarray<*:!numpy.any_dtype> {
// CHECK: return %[[ARG]] : !numpy.ndarray<*:!numpy.any_dtype>
func @no_type_bound(%arg0: !numpy.ndarray<*:!numpy.any_dtype>) -> !numpy.ndarray<*:!numpy.any_dtype> {
return %arg0 : !numpy.ndarray<*:!numpy.any_dtype>
}
// CHECK-LABEL: func @call(
// CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<[2,3,?]:f32>) -> !numpy.ndarray<*:!numpy.any_dtype> {
// CHECK: %[[SHAPE_ERASED:.*]] = numpy.static_info_cast %[[ARG]] : !numpy.ndarray<[2,3,?]:f32> to !numpy.ndarray<*:!numpy.any_dtype>
// CHECK: %[[SHAPED:.*]] = numpy.static_info_cast %[[SHAPE_ERASED]] : !numpy.ndarray<*:!numpy.any_dtype> to !numpy.ndarray<[2,3,?]:f32>
// CHECK: %[[RES:.*]] = call @call(%[[SHAPED]]) : (!numpy.ndarray<[2,3,?]:f32>) -> !numpy.ndarray<*:!numpy.any_dtype>
// CHECK: return %[[SHAPE_ERASED]] : !numpy.ndarray<*:!numpy.any_dtype>
func @call(%arg0: !numpy.ndarray<*:!numpy.any_dtype> {torch.type_bound = !numpy.ndarray<[2,3,?]:f32>}) -> !numpy.ndarray<*:!numpy.any_dtype> {
%0 = call @call(%arg0) : (!numpy.ndarray<*:!numpy.any_dtype>) -> !numpy.ndarray<*:!numpy.any_dtype>
return %arg0 : !numpy.ndarray<*:!numpy.any_dtype>
}
// CHECK-LABEL: func @none_return() {
// CHECK: %[[NONE:.*]] = basicpy.singleton : !basicpy.NoneType
// CHECK: return
func @none_return() -> !basicpy.NoneType {
%1 = basicpy.singleton : !basicpy.NoneType
return %1 : !basicpy.NoneType
}
// CHECK-LABEL: func @none_call_return() {
// CHECK: call @none_return() : () -> ()
// CHECK: %[[NONE:.*]] = basicpy.singleton : !basicpy.NoneType
// CHECK: "test.use"(%[[NONE]]) : (!basicpy.NoneType) -> ()
// CHECK: return
func @none_call_return() {
%0 = call @none_return() : () -> !basicpy.NoneType
"test.use"(%0) : (!basicpy.NoneType) -> ()
return
}