torch-mlir/test/Dialect/Numpy/ops.mlir

53 lines
1.5 KiB
MLIR

// RUN: npcomp-opt -split-input-file %s | npcomp-opt | FileCheck --dump-input=fail %s
// CHECK-LABEL: @any_dtype
func @any_dtype(%arg0: tensor<*x!numpy.any_dtype>) -> (tensor<*x!numpy.any_dtype>) {
return %arg0 : tensor<*x!numpy.any_dtype>
}
// -----
// CHECK-LABEL: @builtin_ufunc
module @builtin_ufunc {
// CHECK: numpy.builtin_ufunc @numpy.add
numpy.builtin_ufunc @numpy.add
// CHECK: numpy.builtin_ufunc @numpy.custom_sub {some_attr = "foobar"}
numpy.builtin_ufunc @numpy.custom_sub { some_attr = "foobar" }
}
// -----
// CHECK-LABEL: @example_generic_ufunc
module @example_generic_ufunc {
// CHECK: numpy.generic_ufunc @numpy.add(
numpy.generic_ufunc @numpy.add (
// CHECK-SAME: overload(%arg0: i32, %arg1: i32) -> i32 {
overload(%arg0: i32, %arg1: i32) -> i32 {
// CHECK: addi
%0 = addi %arg0, %arg1 : i32
numpy.ufunc_return %0 : i32
},
// CHECK: overload(%arg0: f32, %arg1: f32) -> f32 {
overload(%arg0: f32, %arg1: f32) -> f32 {
// CHECK: addf
%0 = addf %arg0, %arg1 : f32
numpy.ufunc_return %0 : f32
}
)
}
// -----
// CHECK-LABEL: @ufunc_apply_ops
module @ufunc_apply_ops {
numpy.generic_ufunc @numpy.add (
overload(%arg0: i32, %arg1: i32) -> i32 {
%0 = addi %arg0, %arg1 : i32
numpy.ufunc_return %0 : i32
}
)
func @example(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> tensor<*xi32> {
%0 = numpy.ufunc_call @numpy.add(%arg0, %arg1) : (tensor<*xi32>, tensor<*xi32>)
-> tensor<*xi32>
return %0 : tensor<*xi32>
}
}