diff --git a/src/dialect/onnx/onnx.td b/src/dialect/onnx/onnx.td index 29733f7..710f3af 100644 --- a/src/dialect/onnx/onnx.td +++ b/src/dialect/onnx/onnx.td @@ -104,7 +104,7 @@ def ONNXGemmNoBiasOp: ONNX_Op<"GemmNoBias", } def ONNXConvNoBiasOp:ONNX_Op<"ConvNoBias", - [NoSideEffect]> { + [NoSideEffect, DeclareOpInterfaceMethods]> { let summary = "ONNX Conv operation with no Bias operand."; let description = [{ "The convolution operator consumes an input tensor and a filter, and" @@ -112,6 +112,8 @@ def ONNXConvNoBiasOp:ONNX_Op<"ConvNoBias", }]; let arguments = (ins AnyTypeOf<[AnyMemRef, AnyTensor]>:$X, AnyTypeOf<[AnyMemRef, AnyTensor]>:$W); let results = (outs AnyTypeOf<[AnyMemRef, AnyTensor]>); + + let verifier = [{ return ::verify(*this); }]; } def ONNXMaxPoolSingleOutOp: ONNX_Op<"MaxPoolSingleOut", diff --git a/src/dialect/onnx/onnx_ops.cpp b/src/dialect/onnx/onnx_ops.cpp index e2e60c3..b36448d 100644 --- a/src/dialect/onnx/onnx_ops.cpp +++ b/src/dialect/onnx/onnx_ops.cpp @@ -489,7 +489,7 @@ void ONNXReshapeOp::inferShapes() { void ONNXTransposeOp::inferShapes() { // Cannot infer shape if no shape exists. if (!getOperand().getType().isa()) - emitError("Shape tensor not ranked."); + return; // Naive transposition which handles the default case of // reversing the shape of the tensor (similar to numpy.transpose). @@ -525,6 +525,181 @@ LogicalResult verify(ONNXTransposeOp op) { return success(); } +//===----------------------------------------------------------------------===// + +// Conv + +// For this operation, we define the attributes once in the original Conv +// operation class. There is no need to redefine the attribute names for the +// other classes based on Conv. +void ONNXConvNoBiasOp::inferShapes() { + // Generic shape for data input X and weight tensor W: + // X: (N x C x D1 x D2 ... x Dn) + // W: (M x C/group x k1 x k2 x ... x kn) + + // Cannot infer shape if no shape exists. + if (!getOperand(0).getType().isa() || + !getOperand(1).getType().isa()) + return; + + auto dataTy = getOperand(0).getType().cast(); + auto weightTy = getOperand(1).getType().cast(); + auto dataShape = dataTy.getShape(); + auto weightShape = weightTy.getShape(); + + // Check that shape of weight and data have same length. + if (dataShape.size() != weightShape.size()) + emitError("Weight size not compatible with data size."); + + // Required attribute auto_pad defaults to NOTSET. + auto autoPad = getAttrOfType( + ONNXConvOp::getAutoPadAttrName()).getValue(); + // Group is a required attribute and should have default value of 1. + int64_t group = getAttrOfType( + ONNXConvOp::getGroupAttrName()).getInt(); + // Check that the X.shape[1] == (W.shape[1] * group) == C condition holds. + if (dataShape[1] != (weightShape[1] * group)) + emitError("Channel dimension mismatch."); + + // Note: the value of the group attribut only impacts the way the + // computation is carried out and not the actual output size. + + // First two output dimensions consist of the number of batches and the + // number of kernels being applied. + // + SmallVector dims; + // Insert batch size. + dims.emplace_back(dataShape[0]); + // Insert number of filters being applied (number of output channels). + dims.emplace_back(weightShape[0]); + + // Spatial dimensions of the output are computed using the formula: + // + // dim = (inputDim - kernelDim + startPadding + endPadding) / stride + 1 + // + SmallVector outSpatialDims; + // Number of spatial dimensions. + int32_t nDims = dataShape.size() - 2; + + // Initialize dimenions based on the input spatial dimensions. + for (int i = 2; i < dataShape.size(); ++i) + outSpatialDims.emplace_back(dataShape[i]); + + // Use kernel_shape attribute if present otherwise use size from weight + // argument. + SmallVector kernelDims; + if (auto kernelShape = getAttrOfType( + ONNXConvOp::getKernelShapeAttrName())) { + if (kernelShape.getValue().size() != nDims) + emitError("kernel_shape length incompatible with spatial dimensions."); + for (int i = 0; i < nDims; ++i) + kernelDims.emplace_back( + (kernelShape.getValue()[i]).cast().getInt()); + } else { + for (int i = 0; i < nDims; ++i) + kernelDims.emplace_back(weightShape[i + 2]); + } + + // Check if dilations attribute is present. + // If it is then compute new kernel size that includes the receptive field. + // In this calculation we assume that the receptive field pixels must all be + // within the bounds of the image. In this case the new kernel size is given + // by: + // + // ( K + 1 ) * d - 1 + // where K is a kernel dimension and d is the dilation along that axis. + // + // From a dimensionality perspective the kernel size becomes the dilated + // kernel size. + if (auto dilations = getAttrOfType( + ONNXConvOp::getDilationsAttrName())) { + if (dilations.getValue().size() != nDims) + emitError("dilations length incompatible with spatial dimensions."); + for (int i = 0; i < nDims; ++i) + kernelDims[i] = (kernelDims[i] + 1) * + (dilations.getValue()[i]).cast().getInt() - 1; + } + + // Subtract kernel dimensions from input data dimensions. + for (int i = 0; i < nDims; ++i) + outSpatialDims[i] -= kernelDims[i]; + + // Add padding information. + if (autoPad == "NOTSET") { + // Use pads to to determine the padding. If attribute is not + // present then pads is considered to be all zeros (no padding). + if (auto pads = getAttrOfType( + ONNXConvOp::getPadsAttrName())) { + // pads consists of two entries for each spatial axis. + if (pads.getValue().size() != 2 * nDims) + emitError("pads size is not twice the spatial size."); + + for (int i = 0; i < nDims; ++i) { + // Padding for beginning of axis. + int32_t p = (pads.getValue()[i]).cast().getInt(); + outSpatialDims[i] += p; + // Padding for end of axis. + p = (pads.getValue()[i + nDims]).cast().getInt(); + outSpatialDims[i] += p; + } + } + } else if (autoPad == "SAME_UPPER" || autoPad == "SAME_LOWER") { + // Pad input so that output size matches input size. + // Each spatial dimension needs to be padded by a total of: + // + // K - 1 + // + // where K is a kernel spatial dimension. + // Pad as if stride is 1. + for (int i = 0; i < nDims; ++i) + outSpatialDims[i] += kernelDims[i] - 1; + } else if (autoPad == "VALID") { + // No padding + } else { + emitError("Unexpected attribute value for auto_pad."); + } + + // Strides + if (auto strides = getAttrOfType( + ONNXConvOp::getStridesAttrName())) { + if (strides.getValue().size() != nDims) + emitError("strides length incompatible with spatial dimensions."); + for (int i = 0; i < nDims; ++i) { + int64_t stride = + (strides.getValue()[i]).cast().getInt(); + outSpatialDims[i] = floor(outSpatialDims[i] / stride); + } + } + + for (int i = 0; i < nDims; ++i) + outSpatialDims[i] += 1; + + dims.append(outSpatialDims.begin(), outSpatialDims.end()); + getResult().setType(RankedTensorType::get(dims, dataTy.getElementType())); +} + +LogicalResult verify(ONNXConvNoBiasOp op) { + auto module = op.getParentOfType(); + if (!module) + op.emitError("expected to belong to a module"); + + auto autoPadAttr = op.getAttrOfType( + ONNXConvOp::getAutoPadAttrName()); + if (!autoPadAttr) + op.emitError("auto_pad attribute not specified."); + if (autoPadAttr.getValue() != "NOTSET") + if (auto pads = op.getAttrOfType( + ONNXConvOp::getPadsAttrName())) + op.emitError("auto_pad and pads are both set."); + + auto groupAttr = + op.getAttrOfType(ONNXConvOp::getGroupAttrName()); + if (!groupAttr) + op.emitError("group attribute not specified."); + + return success(); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// diff --git a/src/dialect/onnx/onnxop.inc b/src/dialect/onnx/onnxop.inc index e87a01a..3a54fa0 100644 --- a/src/dialect/onnx/onnxop.inc +++ b/src/dialect/onnx/onnxop.inc @@ -324,6 +324,15 @@ def ONNXConvOp:ONNX_Op<"Conv", }]; let arguments = (ins AnyTypeOf<[AnyMemRef, AnyTensor]>:$X, AnyTypeOf<[AnyMemRef, AnyTensor]>:$W, AnyTypeOf<[AnyMemRef, AnyTensor]>:$B); let results = (outs AnyTypeOf<[AnyMemRef, AnyTensor]>); + + let extraClassDeclaration = [{ + static StringRef getAutoPadAttrName() { return "auto_pad"; } + static StringRef getDilationsAttrName() { return "dilations"; } + static StringRef getGroupAttrName() { return "group"; } + static StringRef getKernelShapeAttrName() { return "kernel_shape"; } + static StringRef getPadsAttrName() { return "pads"; } + static StringRef getStridesAttrName() { return "strides"; } + }]; } def ONNXConvIntegerOp:ONNX_Op<"ConvInteger", diff --git a/src/pass/shape_inference_pass.cpp b/src/pass/shape_inference_pass.cpp index 3226f16..5239904 100644 --- a/src/pass/shape_inference_pass.cpp +++ b/src/pass/shape_inference_pass.cpp @@ -117,7 +117,8 @@ public: op->getName().getStringRef() != "onnx.GemmNoBias" && op->getName().getStringRef() != "onnx.Reshape" && op->getName().getStringRef() != "onnx.Transpose" && - op->getName().getStringRef() != "onnx.Softmax") + op->getName().getStringRef() != "onnx.Softmax" && + op->getName().getStringRef() != "onnx.ConvNoBias") return false; return llvm::any_of(op->getResultTypes(), [](Type result_type) { return !result_type.isa(); diff --git a/test/backend/test.py b/test/backend/test.py index 60ca4a8..1b257c8 100644 --- a/test/backend/test.py +++ b/test/backend/test.py @@ -144,8 +144,8 @@ test_to_enable = [ #"test_sum_two_inputs_cpu", <- error # Reciprocal Op: - #"test_reciprocal_cpu", <- error on shape inference. - #"test_reciprocal_example_cpu", <- error on shape inference. + "test_reciprocal_cpu", + "test_reciprocal_example_cpu", ] # Extract name of all test cases. diff --git a/test/mlir/onnx/onnx_shape_inference.mlir b/test/mlir/onnx/onnx_shape_inference.mlir index 7057582..e24332f 100644 --- a/test/mlir/onnx/onnx_shape_inference.mlir +++ b/test/mlir/onnx/onnx_shape_inference.mlir @@ -1,7 +1,10 @@ // RUN: onnf-opt --shape-inference %s -split-input-file | FileCheck %s +//===----------------------------------------------------------------------===// /// Test the default behavior of transpose when no information for the -/// permutation of the axes is provided. +/// permutation of the axes is provided and when a permutation is provided. +//===----------------------------------------------------------------------===// + func @test_default_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> { %0 = "onnx.Transpose"(%arg0) : (tensor<5x5x1x32xf32>) -> tensor<*xf32> "std.return"(%0) : (tensor<*xf32>) -> () @@ -12,6 +15,7 @@ func @test_default_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> { } /// Test shape inference for transposition when perm attribute is specified. + func @test_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> { %0 = "onnx.Transpose"(%arg0) {perm = [2, 0, 3, 1]} : (tensor<5x5x1x32xf32>) -> tensor<*xf32> "std.return"(%0) : (tensor<*xf32>) -> () @@ -112,3 +116,128 @@ func @test_matmul_8(%arg0 : tensor<32x64xf32>, %arg1 : tensor<64x128xf32>) -> te // CHECK: [[RES8:%.+]] = "onnx.MatMul"(%arg0, %arg1) : (tensor<32x64xf32>, tensor<64x128xf32>) -> tensor<32x128xf32> // CHECK: return [[RES8]] : tensor<32x128xf32> } + +//===----------------------------------------------------------------------===// +/// Test shape inference for ConvNoBias operation and all its attributes. +//===----------------------------------------------------------------------===// + +/// Default and required attributes. + +func @test_conv_no_bias_1(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_1 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x27x58xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x27x58xf32> + +/// kernel_shape attribute. + +func @test_conv_no_bias_2(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, kernel_shape = [8, 9]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_2 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, kernel_shape = [8, 9]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x25x56xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x25x56xf32> + +/// pads attribute. +/// Use pads to make output size equal to input size by adding K - 1 to the result. + +func @test_conv_no_bias_3(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x10xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, pads = [2, 4, 3, 5]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_3 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, pads = [2, 4, 3, 5]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<1x5x32x64xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x32x64xf32> + +/// auto_pad set to SAME_UPPER and SAME_LOWER. + +func @test_conv_no_bias_4(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x10xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_UPPER", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_4 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_UPPER", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<1x5x32x64xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x32x64xf32> + +func @test_conv_no_bias_5(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x10xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_LOWER", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_5 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_LOWER", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<1x5x32x64xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x32x64xf32> + +/// auto_pad set to VALID. + +func @test_conv_no_bias_6(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x10xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "VALID", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_6 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "VALID", group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x10xf32>) -> tensor<1x5x27x55xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x27x55xf32> + +/// With strides attribute. + +func @test_conv_no_bias_7(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, strides = [2, 3]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_7 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, strides = [2, 3]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x14x20xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x14x20xf32> + +/// auto_pad set to SAME_UPPER with strides attribute. +/// The auto_pad will pas as if stride is equal to 1. + +func @test_conv_no_bias_8(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_UPPER", group = 1 : i32, strides = [2, 3]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_8 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_UPPER", group = 1 : i32, strides = [2, 3]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x16x22xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x16x22xf32> + +/// dilations attribute. + +func @test_conv_no_bias_9(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, dilations = [2, 3]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_9 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", dilations = [2, 3], group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x20x42xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x20x42xf32> + +/// dilations attribute with stride. + +func @test_conv_no_bias_10(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", group = 1 : i32, dilations = [2, 3], strides = [2, 2]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_10 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "NOTSET", dilations = [2, 3], group = 1 : i32, strides = [2, 2]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x10x21xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x10x21xf32> + +/// dilations attribute with auto_pad set to SAME_UPPER. + +func @test_conv_no_bias_11(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> { + %0 = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_UPPER", group = 1 : i32, dilations = [2, 3]} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_conv_no_bias_11 +// CHECK: [[RES_ATTR:%.+]] = "onnx.ConvNoBias"(%arg0, %arg1) {auto_pad = "SAME_UPPER", dilations = [2, 3], group = 1 : i32} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>) -> tensor<1x5x32x64xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x32x64xf32>