From 51ec8d1130c4ed61785750d753d0470521ef3134 Mon Sep 17 00:00:00 2001 From: Gunhyun Park Date: Tue, 23 Apr 2024 13:22:24 -0700 Subject: [PATCH] Update conv docs and tests per interpreter guideline (#2239) This is part 3 of #1964 to implement the remaining parts of #1314. One notable change in TypeInference.cpp is (C27), whose verification differs whether element type is quantized. We have the following constraints in the spec (excluding quantization-related constraints C28-C33): ``` (I1) `lhs` tensor. (I2) `rhs` tensor. (I3) `window_strides` 1-dimensional tensor constant of type `si64`. (I4) `padding` 2-dimensional tensor constant of type `si64`. (I5) `lhs_dilation` 1-dimensional tensor constant of type `si64`. (I6) `rhs_dilation` 1-dimensional tensor constant of type `si64`. (I7) `window_reversal` 1-dimensional tensor constant of type `i1`. (I8) `input_batch_dimension` constant of type `si64`. (I9) `input_feature_dimension` constant of type `si64`. (I10) `input_spatial_dimensions` 1-dimensional tensor constant of type `si64`. (I11) `kernel_input_feature_dimension` constant of type `si64`. (I12) `kernel_output_feature_dimension` constant of type `si64`. (I13) `kernel_spatial_dimensions` 1-dimensional tensor constant of type `si64`. (I14) `output_batch_dimension` constant of type `si64`. (I15) `output_feature_dimension` constant of type `si64`. (I16) `output_spatial_dimensions` 1-dimensional tensor constant of type `si64`. (I17) `feature_group_count` constant of type `si64`. (I18) `batch_group_count` constant of type `si64`. (I19) `precision_config` variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST`. (C1) `N = rank(lhs) = rank(rhs)`. (C2) `size(window_strides) = N - 2`. (C3) `0 < window_strides`. (C4) `shape(padding) = [N - 2, 2]`. (C5) `size(lhs_dilation) = N - 2`. (C6) `0 < lhs_dilation`. (C7) `size(rhs_dilation) = N - 2`. (C8) `0 < rhs_dilation`. (C9) `size(window_reversal) = N - 2`. (C10) `dim(lhs, input_batch_dimension) % batch_group_count = 0`. (C11) `dim(lhs, input_feature_dimension) % feature_group_count = 0`. (C12) `size(input_spatial_dimensions) = N - 2`. (C13) Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * `is_unique(input_dimensions)`. * `0 <= input_dimensions < N`. (C14) `dim(rhs, kernel_input_feature_dimension = dim(lhs, input_feature_dimension) / feature_group_count`. (C15) `dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0`. (C16) `dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0`. (C17) `size(kernel_spatial_dimensions) = N - 2`. (C18) Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * `is_unique(kernel_dimensions)`. * `0 <= kernel_dimensions < N`. (C19) `size(output_spatial_dimensions) = N - 2`. (C20) Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * `is_unique(output_dimensions)`. * `0 <= output_dimensions < N`. (C21) `0 < feature_group_count`. (C22) `0 < batch_group_count`. (C23) `feature_group_count = 1 or batch_group_count = 1`. (C24) `size(precision_config) = 2`. (C25) `dim(result, result_dim)` is defined as: * `dim(lhs, input_batch_dimension) / batch_group_count` if `result_dim = output_batch_dimension`. * `dim(rhs, kernel_output_feature_dimension)` if `result_dim = output_feature_dimension`. * `num_windows` otherwise, where: * `output_spatial_dimensions[spatial_dim] = result_dim`. * `lhs_dim = input_spatial_dimensions[spatial_dim]`. * `rhs_dim = kernel_spatial_dimensions[spatial_dim]`. * `dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1`. * `padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]`. * `dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1`. * `is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]`. * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. (C26) `rank(result) = N`. (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. ``` These constraints will be comprehensively covered by the following tests: ``` I1: a) `lhs` tensor. (Covered by ODS). I2: a) `rhs` tensor. (Covered by ODS). I3: a) `window_strides` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`window_strides`) != `si64`. (Covered by ODS). I4: a) `padding` is not a 2-dimensional tensor. b) element_type(`padding`) != `si64`. (Covered by ODS). I5: a) `lhs_dilation` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`lhs_dilation`) != `si64`. (Covered by ODS). I6: a) `rhs_dilation` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`rhs_dilation`) != `si64`. (Covered by ODS). I7: a) `window_reversal` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`window_reversal`) != `i1`. (Covered by ODS). I8: a) element_type(`input_batch_dimension`) != `si64`. (Covered by ODS). I9: a) element_type(`input_feature_dimension`) != `si64`. (Covered by ODS). I10: a) `input_spatial_dimensions` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`input_spatial_dimensions`) != `si64`. (Covered by ODS). I11: a) element_type(`kernel_input_feature_dimension`) != `si64`. (Covered by ODS). I12: a) element_type(`kernel_output_feature_dimension`) != `si64`. (Covered by ODS). I13: a) `kernel_spatial_dimensions` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`kernel_spatial_dimensions`) != `si64`. (Covered by ODS). I14: a) element_type(`output_batch_dimension`) != `si64`. (Covered by ODS). I15: a) element_type(`output_feature_dimension`) != `si64`. (Covered by ODS). I16: a) `output_spatial_dimensions` is not a 1-dimensional tensor. (Covered by ODS). b) element_type(`output_spatial_dimensions`) != `si64`. (Covered by ODS). I17: a) element_type(`feature_group_count`) != `si64`. (Covered by ODS). I18: a) element_type(`batch_group_count`) != `si64`. (Covered by ODS). I19: a) `precision_config` does not have variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST`. (Covered by ODS). C1: a) N = rank(`lhs`) != rank(`rhs`). C2: a) size(`window_strides`) != N - 2. C3: a) `window_strides[i]` <= 0 for any i in [0, size(`window_strides`)). C4: a) dim(`padding`, 0) != N - 2. b) dim(`padding`, 1) != 2. C5: a) size(`lhs_dilation`) != N - 2. C6: a) `lhs_dilation[i]` <= 0 for any i in [0, size(`lhs_dilation`)). C7: a) size(`rhs_dilation`) != N - 2. C8: a) `rhs_dilation[i]` <= 0 for any i in [0, size(`rhs_dilation`)). C9: a) size(`window_reversal`) != N - 2. C10: a) `dim(lhs, input_batch_dimension) % batch_group_count != 0`. C11: a) `dim(lhs, input_feature_dimension) % feature_group_count != 0`. C12: a) size(`input_spatial_dimensions`) != N - 2. C13: a) Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * Any dimensions in `input_dimensions` are not unique. b) Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * For any i in `input_dimensions`, i < 0. c) Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * For any i in `input_dimensions`, i >= N. C14: a) `dim(rhs, kernel_input_feature_dimension != dim(lhs, input_feature_dimension) / feature_group_count`. C15: a) `dim(rhs, kernel_output_feature_dimension) % batch_group_count != 0`. C16: a) `dim(rhs, kernel_output_feature_dimension) % feature_group_count != 0`. C17: a) size(`kernel_spatial_dimensions`) != N - 2. C18: a) Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * Any dimensions in `kernel_dimensions` are not unique. b) Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * For any i in$ `kernel_dimensions`, i < 0. c) Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * For any i in `kernel_dimensions`, i >= N. C19: a) size(`output_spatial_dimensions`) != N - 2. C20: a) Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * Any dimensions in `output_dimensions` are not unique. b) Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * For any i in `output_dimensions`, i < 0. c) Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * For any i in `output_dimensions`, i >= N. C21: a) `feature_group_count <= 0`. C22: a) `batch_group_count <= 0`. C23: a) `feature_group_count` != 1 and `batch_group_count` != 1. C24: a) size(`precision_config`) != 2. C25: a) For result_dim in [0, N): `dim(result, result_dim)` != `dim(lhs, input_batch_dimension) / batch_group_count`, if `result_dim = output_batch_dimension`. b) For result_dim in [0, N): `dim(result, result_dim)` != `dim(rhs, kernel_output_feature_dimension)`, if `result_dim = output_feature_dimension`. c) For result_dim in [0, N): `dim(result, result_dim)` != `num_windows` otherwise, where: * `output_spatial_dimensions[spatial_dim] = result_dim`. * `lhs_dim = input_spatial_dimensions[spatial_dim]`. * `rhs_dim = kernel_spatial_dimensions[spatial_dim]`. * `dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) == 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1`. * `padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]`. * `dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) == 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1`. * `num_windows = (padded_input_shape[lhs_dim] == 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]) ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. C26: a) rank(result) != N. C27: a) element_type(`lhs`) != element_type(`rhs`). ``` If we drop the "Covered by ODS" pieces, this will leave us with the following test cases: ``` I4a: `padding` is not a 2-dimensional tensor. C1a: rank(`lhs`) != rank(`rhs`) != N. C2a: size(`window_strides`) != N - 2. C3a: `window_strides[i]` <= 0 for any i in [0, size(`window_strides`)). C4a: dim(`padding`, 0) != N - 2. C4b: dim(`padding`, 1) != 2. C5a: size(`lhs_dilation`) != N - 2. C6a: `lhs_dilation[i]` <= 0 for any i in [0, size(`lhs_dilation`)). C7a: size(`rhs_dilation`) != N - 2. C8a: `rhs_dilation[i]` <= 0 for any i in [0, size(`rhs_dilation`)). C9a: size(`window_reversal`) != N - 2. C10a: `dim(lhs, input_batch_dimension) % batch_group_count != 0`. C11a: `dim(lhs, input_feature_dimension) % feature_group_count != 0`. C12a: size(`input_spatial_dimensions`) != N - 2. C13a: Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * Any dimensions in `input_dimensions` are not unique. C13b: Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * For any i in `input_dimensions`, i < 0. C13c: Given `input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]`: * For any i in `input_dimensions`, i >= N. C14a: `dim(rhs, kernel_input_feature_dimension != dim(lhs, input_feature_dimension) / feature_group_count`. C15a: `dim(rhs, kernel_output_feature_dimension) % batch_group_count != 0`. C16a: `dim(rhs, kernel_output_feature_dimension) % feature_group_count != 0`. C17a: size(`kernel_spatial_dimensions`) != N - 2. C18a: Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * Any dimensions in `kernel_dimensions` are not unique. C18b: Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * For any i in$ `kernel_dimensions`, i < 0. C18c: Given `kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`: * For any i in `kernel_dimensions`, i >= N. C19a: size(`output_spatial_dimensions`) != N - 2. C20a: Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * Any dimensions in `output_dimensions` are not unique. b) Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * For any i in `output_dimensions`, i < 0. c) Given `output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]`: * For any i in `output_dimensions`, i >= N. C21a: `feature_group_count <= 0`. C22a: `batch_group_count <= 0`. C23a: `feature_group_count` != 1 and `batch_group_count` != 1. C24a: size(`precision_config`) != 2. C25a: For result_dim in [0, N): `dim(result, result_dim)` != `dim(lhs, input_batch_dimension) / batch_group_count`, if `result_dim = output_batch_dimension`. C25b: For result_dim in [0, N): `dim(result, result_dim)` != `dim(rhs, kernel_output_feature_dimension)`, if `result_dim = output_feature_dimension`. C25c: For result_dim in [0, N): `dim(result, result_dim)` != `num_windows` otherwise, where: * `output_spatial_dimensions[spatial_dim] = result_dim`. * `lhs_dim = input_spatial_dimensions[spatial_dim]`. * `rhs_dim = kernel_spatial_dimensions[spatial_dim]`. * `dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) == 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1`. * `padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]`. * `dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) == 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1`. * `num_windows = (padded_input_shape[lhs_dim] == 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]) ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. C26a: rank(result) != N. C27a: element_type(`lhs`) != element_type(`rhs`). ``` Notes: * (new C24) is left untouched as there are still pending action item regarding the number of precision config values allowed in #879. closes #2092 --- stablehlo/dialect/StablehloOps.td | 36 +- stablehlo/dialect/TypeInference.cpp | 176 +++-- stablehlo/tests/infer_stablehlo.mlir | 200 +++++- stablehlo/tests/ops_stablehlo.mlir | 24 +- stablehlo/tests/verify_conv.mlir | 940 +++++++++++++-------------- 5 files changed, 773 insertions(+), 603 deletions(-) diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td index c0236e0639a..a4d2ec4e4e9 100644 --- a/stablehlo/dialect/StablehloOps.td +++ b/stablehlo/dialect/StablehloOps.td @@ -2216,24 +2216,32 @@ def StableHLO_ConvolutionOp : StableHLO_Op<"convolution", Example: ```mlir - %result = "stablehlo.convolution"(%lhs, %rhs) { - window_strides = array, - padding = dense<0> : tensor<2x2xi64>, - lhs_dilation = array, - rhs_dilation = array, - window_reversal = array, - dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, - feature_group_count = 1 : i64, - batch_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32> + %result = stablehlo.convolution(%lhs, %rhs) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = { + stride = [4, 4], + pad = [[0, 0], [0, 0]], + lhs_dilate = [2, 2], + rhs_dilate = [1, 1], + reverse = [0, 0] + } { + feature_group_count = 1 : i64, + batch_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64> ``` }]; let arguments = !con( (ins - HLO_Tensor:$lhs, - HLO_TensorOrPerAxisQuantizedTensor:$rhs), - StableHLO_ConvolutionAttributes.attributes); + HLO_Tensor:$lhs, /*convolution_i1*/ + HLO_TensorOrPerAxisQuantizedTensor:$rhs), /*convolution_i2*/ + StableHLO_ConvolutionAttributes.attributes /*convolution_i3, convolution_i4, + convolution_i5, convolution_i6, convolution_i7, convolution_i8, + convolution_i9, convolution_i10, convolution_i11, convolution_i12, + convolution_i13, convolution_i14, convolution_i15, convolution_i16, + convolution_i17, convolution_i18, convolution_i19*/ + ); let results = (outs HLO_TensorOrPerAxisQuantizedTensor); let hasVerifier = 1; diff --git a/stablehlo/dialect/TypeInference.cpp b/stablehlo/dialect/TypeInference.cpp index 42c181e4415..c3b5381ef6b 100644 --- a/stablehlo/dialect/TypeInference.cpp +++ b/stablehlo/dialect/TypeInference.cpp @@ -598,25 +598,30 @@ verifyWindowAttributesAndInferWindowDimensions( " to have same dimension-size as size of window dimensions (", windowDimensions.size(), "), but got: ", attrSize, "."); }; - // reduce_window_c6, select_and_scatter_c6 + // convolution_c2, reduce_window_c6, select_and_scatter_c6 if (failed(verifySize(windowStrides.size(), "window-strides"))) return failure(); - // reduce_window_c8 + + // convolution_c5, reduce_window_c8 if (failed(verifySize(lhsDilation.size(), "base-dilation factors"))) return failure(); - // reduce_window_c10 + + // convolution_c7, reduce_window_c10 if (failed(verifySize(rhsDilation.size(), "window-dilation factors"))) return failure(); - // reduce_window_c12 + + // convolution_c4, reduce_window_c12 if (failed(verifySize(padding.size(), "padding-entries"))) return failure(); + + // convolution_c9 if (failed(verifySize(windowReversal.size(), "window-reversal"))) return failure(); SmallVector window(windowDimensions.size()); for (size_t i = 0; i < windowDimensions.size(); i++) { WindowDimension& dim = window[i]; - dim.size = windowDimensions[i]; + // reduce_window_c5, select_and_scatter_c5 if (!isDynamicDimSize(dim.size) && dim.size <= 0) return emitOptionalError(loc, @@ -624,21 +629,24 @@ verifyWindowAttributesAndInferWindowDimensions( "-th window dimension, but got ", dim.size, "."); if (!windowStrides.empty()) dim.stride = windowStrides[i]; - // reduce_window_c7, select_and_scatter_c7 + + // convolution_c3, reduce_window_c7, select_and_scatter_c7 if (dim.stride <= 0) return emitOptionalError( loc, "expects window to have positive stride for ", i, "-th window dimension, but got ", dim.stride, "."); if (!lhsDilation.empty()) dim.baseDilation = lhsDilation[i]; - // reduce_window_c9 + + // convolution_c6, reduce_window_c9 if (dim.baseDilation <= 0) return emitOptionalError( loc, "expects window to have positive base dilation factor for ", i, "-th window dimension, but got ", dim.baseDilation, "."); if (!rhsDilation.empty()) dim.windowDilation = rhsDilation[i]; - // reduce_window_c11 + + // convolution_c8, reduce_window_c11 if (dim.windowDilation <= 0) return emitOptionalError( loc, "expects window to have positive window dilation factor for ", i, @@ -987,12 +995,6 @@ LogicalResult verifyRegionNotEmpty(std::optional location, return success(); } -// Checks: -// P1. Same sizes for input, kernel and output spatialDims. -// P2. Spatial and non-spatial dimensions (for input,kernel, &output) should -// be unique and in range [0, num_dims), where num_dims = rank of input -// (lhs/rhs) tensors. -// // Note that the spatial + non-spatial dimensions may not cover all the // dimensions in the range [0,num) because of the presence of 'unknown' // dimensions (ref. `printConvolutionDimensions()`) @@ -1004,7 +1006,7 @@ LogicalResult isSpatialDimensionsValid( int64_t outputFeatureDimension, ArrayRef outputSpatialDimensions, std::optional location) { uint64_t spatialDimNum = inputSpatialDimensions.size(); - // P1. + // convolution_c17, convolution_c19 if ((spatialDimNum != kernelSpatialDimensions.size()) || (spatialDimNum != outputSpatialDimensions.size())) return emitOptionalError(location, @@ -1014,7 +1016,6 @@ LogicalResult isSpatialDimensionsValid( kernelSpatialDimensions.size(), ", and ", outputSpatialDimensions.size(), " resp."); - // P2. SmallVector inputDimNums(spatialDimNum + 2); inputDimNums[0] = inputBatchDimension; inputDimNums[1] = inputFeatureDimension; @@ -1027,37 +1028,40 @@ LogicalResult isSpatialDimensionsValid( std::copy(kernelSpatialDimensions.begin(), kernelSpatialDimensions.end(), windowDimNums.begin() + 2); - SmallVector OutputDimNums(spatialDimNum + 2); - OutputDimNums[0] = outputBatchDimension; - OutputDimNums[1] = outputFeatureDimension; + SmallVector outputDimNums(spatialDimNum + 2); + outputDimNums[0] = outputBatchDimension; + outputDimNums[1] = outputFeatureDimension; std::copy(outputSpatialDimensions.begin(), outputSpatialDimensions.end(), - OutputDimNums.begin() + 2); + outputDimNums.begin() + 2); auto numDims = cast(lhsType).getRank(); const auto inRange = [numDims](int64_t i) { return 0 <= i && i < numDims; }; - + // convolution_c13, convolution_c18, convolution_c20 if (!llvm::all_of(inputDimNums, inRange) || !llvm::all_of(windowDimNums, inRange) || - !llvm::all_of(OutputDimNums, inRange)) + !llvm::all_of(outputDimNums, inRange)) return emitOptionalError(location, "expects input, kernel, and output " "dimension-numbers to be in-range [0, ", numDims, ")."); + // convolution_c13 if (!isUnique(inputDimNums)) return emitOptionalError( location, "expects input dimension-numbers to be unique, got {", inputDimNums, "}."); + // convolution_c18 if (!isUnique(windowDimNums)) return emitOptionalError( location, "expects kernel dimension-numbers to be unique, got {", windowDimNums, "}."); - if (!isUnique(OutputDimNums)) + // convolution_c20 + if (!isUnique(outputDimNums)) return emitOptionalError( location, "expects output dimension-numbers to be unique, got {", - OutputDimNums, "}."); + outputDimNums, "}."); return success(); } @@ -1075,27 +1079,6 @@ LogicalResult verifyPrecisionConfig(std::optional loc, "<= 2 elements."); } -// Verifies the following properties: -// P1. The input, kernel, and output spatial-dimensions are valid. -// P2. Given, -// input-dimensions: b * input-spatial-dims * f -// kernel-dimensions: kernel-spatial-dims * i * o -// output-dimensions: b' * out-spatial-dims * f' -// where b = input-batch-dim -// where f = input-feature-dim -// where i = kernel-input-feature-dim -// where o = kernel-output-feature-dim -// where b' = output-batch-dim -// where f' = output-feature-dim -// Check the following properties w.r.t feature_group_count (fgc) and -// batch_group_count (bgc). -// * fgc > 0, bgc > 0 and !(fgc > 1 && bgc > 1) -// * dim(lhs, b) % bgc == 0 -// * dim(lhs, f) % fgc == 0 and -// dim(lhs, f) / fgc = dim(rhs, i) -// * dim(rhs, o) (or dim(output, f')) % bgc == 0 and -// dim(rhs, o) (or dim(output, f')) % fgc == 0 -// P3. Precision config is null, of size 0 or of size 2. LogicalResult verifyConvolutionAttributes( std::optional location, Type lhsType, Type rhsType, int64_t inputBatchDimension, int64_t inputFeatureDimension, @@ -1105,7 +1088,6 @@ LogicalResult verifyConvolutionAttributes( int64_t outputFeatureDimension, ArrayRef outputSpatialDimensions, int64_t featureGroupCount, int64_t batchGroupCount, std::optional precisionConfig) { - // P1. if (failed(isSpatialDimensionsValid( lhsType, inputBatchDimension, inputFeatureDimension, inputSpatialDimensions, kernelInputFeatureDimension, @@ -1114,17 +1096,19 @@ LogicalResult verifyConvolutionAttributes( location))) return failure(); - // P2. + // convolution_c21 if (featureGroupCount <= 0) return emitOptionalError( location, "expects feature_group_count to be a positive number, got ", featureGroupCount, "."); + // convolution_c22 if (batchGroupCount <= 0) return emitOptionalError( location, "expects batch_group_count to be a positive number, got ", batchGroupCount, "."); + // convolution_c23 if (batchGroupCount > 1 && featureGroupCount > 1) return emitOptionalError( location, @@ -1142,24 +1126,16 @@ LogicalResult verifyConvolutionAttributes( const int64_t kernelOutputFeatures = rankedRhsType.getShape()[kernelOutputFeatureDimension]; - if (!isDynamicDimSize(kernelOutputFeatures)) { - if (kernelOutputFeatures % batchGroupCount != 0) - return emitOptionalError( - location, "expects output feature dimension size (", - kernelOutputFeatures, - ") to be a multiple of batch_group_count. Got batch_group_count = ", - batchGroupCount, "."); - - if (kernelOutputFeatures % featureGroupCount != 0) - return emitOptionalError(location, - "expects kernel output feature dimension (", - kernelOutputFeatures, - ") to be divisible by feature_group_count. For " - "feature_group_count = ", - featureGroupCount, "."); - } + // convolution_c10 + if (!isDynamicDimSize(inputBatch) && inputBatch % batchGroupCount != 0) + return emitOptionalError(location, "expects input batch dimension (", + inputBatch, + ") to be divisible by " + "batch_group_count. Got batch_group_count = ", + batchGroupCount, "."); if (!isDynamicDimSize(inputFeatures)) { + // convolution_c11 if (inputFeatures % featureGroupCount != 0) return emitOptionalError(location, "expects input feature dimension (", inputFeatures, @@ -1167,6 +1143,7 @@ LogicalResult verifyConvolutionAttributes( "feature_group_count = ", featureGroupCount, "."); + // convolution_c14 if (!isDynamicDimSize(kernelInputFeatures) && inputFeatures / featureGroupCount != kernelInputFeatures) return emitOptionalError( @@ -1177,14 +1154,26 @@ LogicalResult verifyConvolutionAttributes( "). Got feature_group_count = ", featureGroupCount, "."); } - if (!isDynamicDimSize(inputBatch) && inputBatch % batchGroupCount != 0) - return emitOptionalError(location, "expects input batch dimension (", - inputBatch, - ") to be divisible by " - "batch_group_count. Got batch_group_count = ", - batchGroupCount, "."); + if (!isDynamicDimSize(kernelOutputFeatures)) { + // convolution_c15 + if (kernelOutputFeatures % batchGroupCount != 0) + return emitOptionalError( + location, "expects output feature dimension size (", + kernelOutputFeatures, + ") to be a multiple of batch_group_count. Got batch_group_count = ", + batchGroupCount, "."); - // P3. + // convolution_c16 + if (kernelOutputFeatures % featureGroupCount != 0) + return emitOptionalError(location, + "expects kernel output feature dimension (", + kernelOutputFeatures, + ") to be divisible by feature_group_count. For " + "feature_group_count = ", + featureGroupCount, "."); + } + + // convolution_c24 if (failed(verifyPrecisionConfig(location, precisionConfig))) return failure(); @@ -1931,15 +1920,6 @@ LogicalResult inferCollectivePermuteOp( return success(); } -/* - * We intend to verify the following properties - * P1. Verify the input, kernel types. - * P2. Verify the convolution atributes. - * P3. Verify and collect the window atributes. - * P4. Verify precision_config attribute. - * P5. Verify the return shape. - * TODO(b/232574102): Verify the element-type of return-value. - */ LogicalResult inferConvolutionOp( std::optional location, Type lhsType, Type rhsType, std::optional> windowStrides, @@ -1957,19 +1937,30 @@ LogicalResult inferConvolutionOp( auto rankedLhsType = cast(lhsType); auto rankedRhsType = cast(rhsType); - // P1. + // convolution_c13 int numDims = rankedLhsType.getRank(); + if (numDims < 2) + return emitOptionalError( + location, + "expects convolution arguments to have >= 2 dimensions. Got: ", + rankedLhsType, " and ", rankedRhsType, "."); + + // convolution_c1 if (numDims != rankedRhsType.getRank()) return emitOptionalError(location, "expects convolution arguments to have same " "number of dimensions. Got: ", rankedLhsType, " and ", rankedRhsType, "."); - if (numDims < 2) + + // convolution_c27 + if (!anyQuantized({rankedLhsType, rankedRhsType}) && + !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), + rankedRhsType.getElementType())) return emitOptionalError( - location, - "expects convolution arguments to have >= 2 dimensions. Got: ", - rankedLhsType, " and ", rankedRhsType, "."); - // P2. + location, "expects lhs and rhs to have compatible element type. Got: ", + rankedLhsType.getElementType(), " and ", + rankedRhsType.getElementType()); + if (failed(verifyConvolutionAttributes( location, lhsType, rhsType, inputBatchDimension, inputFeatureDimension, inputSpatialDimensions, @@ -1979,16 +1970,17 @@ LogicalResult inferConvolutionOp( precisionConfig))) return failure(); + // convolution_c12 if ((size_t)numDims != inputSpatialDimensions.size() + 2) return emitOptionalError(location, "expects convolution arguments to have ", inputSpatialDimensions.size() + 2, " dimensions. Got: ", numDims); - // P3. SmallVector windowDimensions(kernelSpatialDimensions.size()); for (size_t i = 0; i < windowDimensions.size(); i++) windowDimensions[i] = rankedRhsType.getShape()[kernelSpatialDimensions[i]]; + // convolution_c4, convolution_i4 auto paddingOrErr = convertPaddingAttribute(padding, location); if (failed(paddingOrErr)) return failure(); @@ -2001,15 +1993,10 @@ LogicalResult inferConvolutionOp( windowReversal.value_or(ArrayRef{}), location); if (failed(windowOrErr)) return failure(); - // P3. - if (failed(verifyPrecisionConfig(location, precisionConfig))) - return failure(); - - // P5. + // convolution_c25, convolution_c26 SmallVector outputDimensions(rankedLhsType.getShape().size(), ShapedType::kDynamic); - // Infer the output spatial dimensions. auto numSpatialDims = inputSpatialDimensions.size(); SmallVector inputSpatialDimVals(numSpatialDims); for (int64_t i = 0; i < static_cast(numSpatialDims); ++i) @@ -2021,7 +2008,6 @@ LogicalResult inferConvolutionOp( for (int64_t i = 0; i < static_cast(windowOrErr->size()); ++i) outputDimensions[outputSpatialDimensions[i]] = windowOutputShape[i]; - // Infer the output-batch-dimension and output-feature-dimension. const int64_t inputBatch = rankedLhsType.getShape()[inputBatchDimension]; const int64_t kernelOutputFeatures = rankedRhsType.getShape()[kernelOutputFeatureDimension]; @@ -3580,12 +3566,14 @@ LogicalResult verifyConvolutionOp( auto inferredShape = inferredReturnShapes[0]; auto shapedResultType = cast(resultType); + + // convolution_c25 if (failed(verifyCompatibleShape(inferredShape.getDims(), shapedResultType.getShape()))) return emitOptionalError(location, "inferred shape '", dimSizesToString(inferredShape.getDims()), "' ", "is incompatible with return type of operation ", - shapedResultType, ""); + shapedResultType); if (anyQuantized({lhsType, rhsType, resultType})) { return verifyConvolutionOpQuantizationConstraints( diff --git a/stablehlo/tests/infer_stablehlo.mlir b/stablehlo/tests/infer_stablehlo.mlir index f78c7b3d320..c0ca2e78b1a 100644 --- a/stablehlo/tests/infer_stablehlo.mlir +++ b/stablehlo/tests/infer_stablehlo.mlir @@ -160,16 +160,6 @@ func.func @abs(%arg0: tensor<1x2xf32>) -> tensor<1x2xindex> { // ----- -// CHECK-LABEL: @concat -func.func @concat(%arg0: tensor<1xi32>, %arg1: tensor<2xi32>) -> tensor<3xindex> { - %0 = "stablehlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<1xi32>, tensor<2xi32>) -> tensor<3xi32> - // CHECK: types0 = tensor<3xi32> - %1 = "hlo_test_infer.get_return_types"(%0) : (tensor<3xi32>) -> tensor<3xindex> - func.return %1 : tensor<3xindex> -} - -// ----- - // CHECK-LABEL: func @collective_broadcast_c3 func.func @collective_broadcast_c3(%arg0: tensor<16x8xf32>) -> tensor<16x8xindex> { %0 = "stablehlo.collective_broadcast"(%arg0) { @@ -195,6 +185,196 @@ func.func @collective_permute_c5(%arg0: tensor<2x2xi64>) -> tensor<2x2xindex> { // ----- +// CHECK-LABEL: @concat +func.func @concat(%arg0: tensor<1xi32>, %arg1: tensor<2xi32>) -> tensor<3xindex> { + %0 = "stablehlo.concatenate"(%arg0, %arg1) { dimension = 0 : i64 } : (tensor<1xi32>, tensor<2xi32>) -> tensor<3xi32> + // CHECK: types0 = tensor<3xi32> + %1 = "hlo_test_infer.get_return_types"(%0) : (tensor<3xi32>) -> tensor<3xindex> + func.return %1 : tensor<3xindex> +} + +// ----- + +// Invalid rank of output-type. + +func.func @invalid_conv_return_type(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x16xf32> { + // expected-error @+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x8x16xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x16xf32> + func.return %0 : tensor<1x8x16xf32> +} + +// ----- + +// Invalid batch dimension in output-type. Should be equal to +// input-batch-dimension / batch_group_count. + +func.func @invalid_conv_return_type(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<2x8x8x16xf32> { + // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<2x8x8x16xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<2x8x8x16xf32> + func.return %0 : tensor<2x8x8x16xf32> +} + +// ----- + +// Invalid feature dimension in output-type. Should be equal to +// kernel_output_feature_dimension. + +func.func @invalid_conv_return_type(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x32xf32> { + // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x8x8x32xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x32xf32> + func.return %0 : tensor<1x8x8x32xf32> +} + +// ----- + +// The following tests checks the inferred output-type of ConvolutionOp. We +// deliberately put an invalid output-type in these tests so that the +// inffered-type can be highlighted in the error message. + +// Dynamic input-batch-dimension +func.func @invalid_conv_dynamic_shapes(%arg0: tensor, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> { + // expected-error@+1 {{inferred shape '[?, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor, tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> + func.return %0 : tensor<1x1x1x1xf32> +} + +// ----- + +// Dynamic input-feature-dimension: No effect on output dimensions. +func.func @invalid_conv_dynamic_shapes(%arg0: tensor<1x8x8x?xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> { + // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x?xf32>, tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> + func.return %0 : tensor<1x1x1x1xf32> +} + +// ----- + +// Dynamic input-spatial-dimension +func.func @invalid_conv_dynamic_shapes(%arg0: tensor<1x?x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> { + // expected-error@+1 {{inferred shape '[1, ?, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x?x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> + func.return %0 : tensor<1x1x1x1xf32> +} + +// ----- + +// Dynamic kernel-input-feature-dimension: No effect on output dimensions. +func.func @invalid_conv_dynamic_shapes(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x?x16xf32>) -> tensor<1x1x1x1xf32> { + // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x3x?x16xf32>) -> tensor<1x1x1x1xf32> + func.return %0 : tensor<1x1x1x1xf32> +} + +// ----- + +// Dynamic kernel-output-feature-dimension +func.func @check_inferred_type_with_dynamic_input_dims(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x?xf32>) -> tensor<1x1x1x1xf32> { + // expected-error@+1 {{inferred shape '[1, 8, 8, ?]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x3x207x?xf32>) -> tensor<1x1x1x1xf32> + func.return %0 : tensor<1x1x1x1xf32> +} + +// ----- + +// Dynamic kernel-spatial-dimension +func.func @check_inferred_type_with_dynamic_input_dims(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x?x207x16xf32>) -> tensor<1x1x1x1xf32> { + // expected-error@+1 {{inferred shape '[1, 8, ?, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x?x207x16xf32>) -> tensor<1x1x1x1xf32> + func.return %0 : tensor<1x1x1x1xf32> +} + +// ----- + // CHECK-LABEL: @gather_c13 func.func @gather_c13(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>) -> tensor<1x5x8xindex> { %res = "stablehlo.gather"(%operand, %start_indices) { diff --git a/stablehlo/tests/ops_stablehlo.mlir b/stablehlo/tests/ops_stablehlo.mlir index 97b431f951d..1faa8ef5c20 100644 --- a/stablehlo/tests/ops_stablehlo.mlir +++ b/stablehlo/tests/ops_stablehlo.mlir @@ -4777,11 +4777,31 @@ func.func @eltwise_static_and_dynamic_type(%arg0: tensor<10x10xf32>, %arg1: tens // ----- -// CHECK: func @quantized_conv2d +// CHECK-LABEL: func @convolution_operand_element_type_i4 +func.func @convolution_operand_element_type_i4(%arg0: tensor<64x8x8x8xi4>, %arg1: tensor<4x4x8x32xi4>) -> tensor<64x3x3x32xi8> { + // Note: This has been lowered and adapted from: + // %0 = "tf.Conv2D"(%arg0, %arg1) { + // data_format = "NHWC", + // dilations = [1, 2, 2, 1], + // explicit_paddings = [0, 0, 0, 1, 0, 1, 0, 0], + // padding = "EXPLICIT", + // strides = [1, 1, 1, 1]} : + // (tensor<64x8x8x8xf32>, tensor<4x4x8x32xf32>) -> tensor<64x3x3x32xf32> + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[0, 1], [0, 1]], rhs_dilate = [2, 2]} + {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : + (tensor<64x8x8x8xi4>, tensor<4x4x8x32xi4>) -> tensor<64x3x3x32xi8> + func.return %0 : tensor<64x3x3x32xi8> +} + +// ----- + +// CHECK: func @convolution_quantized_conv2d // CHECK: stablehlo.convolution // CHECK-SAME: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f] // CHECK-SAME{LITERAL}: window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} -func.func @quantized_conv2d(%arg0: tensor<1x8x8x207x!quant.uniform>, %arg1: tensor<3x3x207x16x!quant.uniform>) -> tensor<1x8x8x16x!quant.uniform> { +func.func @convolution_quantized_conv2d(%arg0: tensor<1x8x8x207x!quant.uniform>, %arg1: tensor<3x3x207x16x!quant.uniform>) -> tensor<1x8x8x16x!quant.uniform> { %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} diff --git a/stablehlo/tests/verify_conv.mlir b/stablehlo/tests/verify_conv.mlir index b8cdcd0440d..a5a1a68fcbd 100644 --- a/stablehlo/tests/verify_conv.mlir +++ b/stablehlo/tests/verify_conv.mlir @@ -1,9 +1,7 @@ // RUN: stablehlo-opt %s -verify-diagnostics -split-input-file | FileCheck %s -// Valid: Generic convolution - -// CHECK-LABEL: func @main -func.func @main(%arg0 : tensor<100x26x26x32xf32>, %arg1 : tensor<3x3x1x32xf32>) -> +// CHECK-LABEL: func @convolution +func.func @convolution(%arg0 : tensor<100x26x26x32xf32>, %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { %result = "stablehlo.convolution"(%arg0, %arg1) { batch_group_count = 1 : i64, @@ -30,7 +28,28 @@ func.func @main(%arg0 : tensor<100x26x26x32xf32>, %arg1 : tensor<3x3x1x32xf32>) // ----- -// Valid: Test convolution i8xi8 -> i32. +// CHECK: func @convolution_empty_spatial_dimensions +// CHECK: stablehlo.convolution +// CHECK-SAME: dim_numbers = [b, f]x[i, o]->[b, f] +// CHECK-SAME: window = {stride = [], pad = [], lhs_dilate = [], +// CHECK-SAME: rhs_dilate = [], reverse = []} +func.func @convolution_empty_spatial_dimensions(%arg0: tensor<3x2xf16>, + %arg1: tensor<2x2xf16>) -> tuple> { + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, f]x[i, o]->[b, f], + window = {stride = [], pad = [], lhs_dilate = [], rhs_dilate = [], + reverse = []} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } + : (tensor<3x2xf16>, tensor<2x2xf16>) -> tensor<3x2xf16> + %1 = "stablehlo.tuple"(%0) : (tensor<3x2xf16>) -> tuple> + func.return %1 : tuple> +} + +// ----- // CHECK-LABEL: func @convolution_upcast func.func @convolution_upcast(%arg0 : tensor<100x26x26x32xi8>, @@ -59,71 +78,45 @@ func.func @convolution_upcast(%arg0 : tensor<100x26x26x32xi8>, // ----- -// Valid: Empty spatial dimensions - -// CHECK: func @conv_empty_spatial_dimensions -// CHECK: stablehlo.convolution -// CHECK-SAME: dim_numbers = [b, f]x[i, o]->[b, f] -// CHECK-SAME: window = {stride = [], pad = [], lhs_dilate = [], -// CHECK-SAME: rhs_dilate = [], reverse = []} -func.func @conv_empty_spatial_dimensions(%arg0: tensor<3x2xf16>, - %arg1: tensor<2x2xf16>) -> tuple> { +func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { + // expected-error@+3{{Unexpected keyword stide}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, f]x[i, o]->[b, f], - window = {stride = [], pad = [], lhs_dilate = [], rhs_dilate = [], - reverse = []} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } - : (tensor<3x2xf16>, tensor<2x2xf16>) -> tensor<3x2xf16> - %1 = "stablehlo.tuple"(%0) : (tensor<3x2xf16>) -> tuple> - func.return %1 : tuple> + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stide = [2, 1], pad = [[0, 1], [0, 1]], rhs_dilate = [1, 2]} + { batch_group_count = 1 : i64, feature_group_count = 1 : i64} + : (tensor<2x2x3x4xf32>, tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> + func.return %0 : tensor<3x5x5x4xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<2x4x5x2xf32>, - %arg1: tensor<2x2x1x6xf32>) -> tensor<2x3x4x6xf32> { - // expected-error@+1 {{expects input dimension-numbers to be unique, got {0, 0}.}} - %1 = "stablehlo.convolution"(%arg0, %arg1) { - batch_group_count = 1 : i64, - dimension_numbers = #stablehlo.conv, - feature_group_count = 2 : i64, - someattr} : (tensor<2x4x5x2xf32>, tensor<2x2x1x6xf32>) -> - tensor<2x3x4x6xf32> - func.return %1 : tensor<2x3x4x6xf32> +func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { + // expected-error@+3{{expected integer value}} + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [2, b], pad = [[0, 1], [0, 1]], rhs_dilate = [1, 2]} + { batch_group_count = 1 : i64, feature_group_count = 1 : i64} + : (tensor<2x2x3x4xf32>, tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> + func.return %0 : tensor<3x5x5x4xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects convolution arguments to have same number of dimensions. Got: 'tensor<1x8x8x207xf32>' and 'tensor<3x3x207xf32>'.}} +func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { + // expected-error@+3{{Unexpected keyword stride}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [2, 1], pad = [[0, 1], [0, 1]], rhs_dilate = [1, 2], stride=[2,1]} + { batch_group_count = 1 : i64, feature_group_count = 1 : i64} + : (tensor<2x2x3x4xf32>, tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> + func.return %0 : tensor<3x5x5x4xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1xf32>, %arg1: tensor<3xf32>) - -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects convolution arguments to have >= 2 dimensions. Got: 'tensor<1xf32>' and 'tensor<3xf32>'.}} +func.func @convolution_c1(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects convolution arguments to have same number of dimensions. Got: 'tensor<1x8x8x207xf32>' and 'tensor<3x3x207xf32>'.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], @@ -133,35 +126,37 @@ func.func @invalid_conv_dimensions(%arg0: tensor<1xf32>, %arg1: tensor<3xf32>) feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1xf32>, tensor<3xf32>) -> tensor<1x8x8x16xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects the same size for input, kernel and output spatial-dimensions, but got 3, 2, and 2 resp.}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, 2, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> +func.func @convolution_c2(%arg0: tensor<1x4x4x1xi64>, + %arg1: tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi64> { + // expected-error@+1 {{expects lhs and rhs to have compatible element type. Got: 'i64' and 'i32'}} + %0 = "stablehlo.convolution"(%arg0, %arg1) { + window_strides = array, + padding = dense<0> : tensor<2x2xi64>, + lhs_dilation = array, + rhs_dilation = array, + window_reversal = array, + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, + feature_group_count = 1 : i64, + batch_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi64> + func.return %0 : tensor<1x2x2x1xi64> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c3(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects the same size for input, kernel and output spatial-dimensions, but got 2, 3, and 2 resp.}} + // expected-error@+1 {{expects window-strides to have same dimension-size as size of window dimensions (2), but got: 1.}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, 2, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, @@ -174,33 +169,32 @@ func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c4(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects the same size for input, kernel and output spatial-dimensions, but got 2, 2, and 3 resp.}} + // expected-error@+1 {{expects window to have positive stride for 1-th window dimension, but got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, 2, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], + window = {stride = [1, 0], pad = [[1, 1], [1,1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : + precision_config = [#stablehlo, #stablehlo]} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0 : tensor<100x26x26x32xf32>, - %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { - // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} +func.func @convolution_c5(%arg0 : tensor<100x26x26x32xf32>, %arg1 : tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects padding-entries to have same dimension-size as size of window dimensions (2), but got: 3.}} %result = "stablehlo.convolution"(%arg0, %arg1) { batch_group_count = 1 : i64, dimension_numbers = #stablehlo.conv, >, feature_group_count = 1 : i64, lhs_dilation = array, - padding = dense<2> : tensor<2x2xi64>, + padding = dense<2> : tensor<3x2xi64>, rhs_dilation = array, window_strides = array } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> @@ -220,71 +214,34 @@ func.func @invalid_conv_dimensions(%arg0 : tensor<100x26x26x32xf32>, // ----- -func.func @invalid_conv_dimensions(%arg0 : tensor<100x26x26x32xf32>, - %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { - // expected-error@+1 {{expects kernel dimension-numbers to be unique, got {3, 2, 0, 0}.}} - %result = "stablehlo.convolution"(%arg0, %arg1) { - batch_group_count = 1 : i64, - dimension_numbers = #stablehlo.conv, - feature_group_count = 1 : i64, - lhs_dilation = array, - padding = dense<2> : tensor<2x2xi64>, +func.func @convolution_c5(%arg0: tensor<1x4x4x1xi64>, + %arg1: tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64> { + // expected-error@+1 {{expects the shape of padding-attribute to be {N, 2}, but got {2, 3}.}} + %0 = "stablehlo.convolution"(%arg0, %arg1) { + window_strides = array, + padding = dense<0> : tensor<2x3xi64>, + lhs_dilation = array, rhs_dilation = array, - window_strides = array - } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> - tensor<100x28x28x1xf32> - func.return %result : tensor<100x28x28x1xf32> -} - -// ----- - -func.func @invalid_conv_dimensions(%arg0 : tensor<100x26x26x32xf32>, - %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { - // expected-error@+1 {{expects output dimension-numbers to be unique, got {0, 3, 0, 3}.}} - %result = "stablehlo.convolution"(%arg0, %arg1) { - batch_group_count = 1 : i64, - dimension_numbers = #stablehlo.conv, + window_reversal = array, + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, feature_group_count = 1 : i64, - lhs_dilation = array, - padding = dense<2> : tensor<2x2xi64>, - rhs_dilation = array, - window_strides = array - } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> - tensor<100x28x28x1xf32> - func.return %result : tensor<100x28x28x1xf32> + batch_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64> + func.return %0 : tensor<1x2x2x1xi64> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c5(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects batch_group_count to be a positive number, got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], + // expected-error@+1 {{Expected array with 2 elements, got 4 elements instead}} + window = {stride = [1, 1], pad = [[1, 1, 1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { - batch_group_count = 0 : i64, + batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo] } : @@ -294,34 +251,33 @@ func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c6(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects feature_group_count to be a positive number, got 0.}} + // expected-error@+1 {{expects base-dilation factors to have same dimension-size as size of window dimensions (2), but got: 1.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + lhs_dilate = [1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, - feature_group_count = 0 : i64, - precision_config = [#stablehlo, #stablehlo] - } : + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo]} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c7(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects batch_group_count and feature_group_count not to be both greater than 1. Got 2 and 2 resp.}} + // expected-error@+1 {{expects window to have positive base dilation factor for 0-th window dimension, but got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [0, 1], rhs_dilate = [1, 1]} { - batch_group_count = 2 : i64, - feature_group_count = 2 : i64, + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo] } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> @@ -330,79 +286,59 @@ func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c8(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects output feature dimension size (16) to be a multiple of batch_group_count. Got batch_group_count = 3.}} + // expected-error@+1 {{expects window-dilation factors to have same dimension-size as size of window dimensions (2), but got: 1.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + lhs_dilate = [1, 1], rhs_dilate = [1]} { - batch_group_count = 3 : i64, + batch_group_count = 1 : i64, feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : + precision_config = [#stablehlo, #stablehlo]} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects input feature dimension (207) to be a multiple of feature_group_count. Got feature_group_count = 2.}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 2 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> -} - -// ----- - -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects input feature dimension (207) / feature_group_count = kernel input feature dimension (20). Got feature_group_count = 1.}} +func.func @convolution_c9(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects window to have positive window dilation factor for 0-th window dimension, but got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [0, 1]} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x8x8x207xf32>, tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x69x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects kernel output feature dimension (16) to be divisible by feature_group_count. For feature_group_count = 3.}} +func.func @convolution_c10(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects window-reversal to have same dimension-size as size of window dimensions (2), but got: 1.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + lhs_dilate = [1, 1], rhs_dilate = [1, 1], reverse = [false]} { batch_group_count = 1 : i64, - feature_group_count = 3 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x3x69x16xf32>) -> tensor<1x8x8x16xf32> + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo]} : + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<5x8x8x207xf32>, +func.func @convolution_c11(%arg0: tensor<5x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { // expected-error@+1 {{expects input batch dimension (5) to be divisible by batch_group_count. Got batch_group_count = 2.}} %0 = stablehlo.convolution(%arg0, %arg1) @@ -420,117 +356,130 @@ func.func @invalid_conv_dimensions(%arg0: tensor<5x8x8x207xf32>, // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window-strides to have same dimension-size as size of window dimensions (2), but got: 1.}} +func.func @convolution_c12(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects input feature dimension (207) to be a multiple of feature_group_count. Got feature_group_count = 2.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1], pad = [[1, 1], [1, 1]], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, - feature_group_count = 1 : i64, + feature_group_count = 2 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects base-dilation factors to have same dimension-size as size of window dimensions (2), but got: 1.}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> -} - -// ----- - -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window-dilation factors to have same dimension-size as size of window dimensions (2), but got: 1.}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> +// This is an positive test in MLIR-HLO: +// https://github.com/tensorflow/mlir-hlo/blob/master/tests/Dialect/mhlo/ops.mlir#L3829 +// but negative here: stablehlo.convolution does no support unknown dimenstion +// dim_numbers = [b, 0, 1, ?, f]x[0, 1, ?, i, o]->[?, b, 0, 1, f] +// window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} +func.func @convolution_c13(%arg0: tensor<1x8x8x32x207xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<32x1x8x8x16xf32> { + // expected-error@+1{{expects convolution arguments to have 4 dimensions. Got: 5}} + %0 = "stablehlo.convolution"(%arg0, %arg1) {batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, feature_group_count = 1 : i64, lhs_dilation = array, padding = dense<1> : tensor<2x2xi64>, precision_config = [#stablehlo, #stablehlo], rhs_dilation = array, window_strides = array} : + (tensor<1x8x8x32x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<32x1x8x8x16xf32> + func.return %0 : tensor<32x1x8x8x16xf32> } // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window-reversal to have same dimension-size as size of window dimensions (2), but got: 1.}} +func.func @convolution_c14(%arg0: tensor<1xf32>, %arg1: tensor<3xf32>) + -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects convolution arguments to have >= 2 dimensions. Got: 'tensor<1xf32>' and 'tensor<3xf32>'.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1], reverse = [false]} + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1xf32>, tensor<3xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects padding-entries to have same dimension-size as size of window dimensions (2), but got: 1.}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> +func.func @convolution_c14(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input dimension-numbers to be unique, got {0, 0, 1, 2}.}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - // expected-error@+1 {{Expected array with 2 elements, got 4 elements instead}} - window = {stride = [1, 1], pad = [[1, 1, 1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> +func.func @convolution_c14(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> } // ----- -func.func @invalid_conv_dimensions(%arg0 : tensor<100x26x26x32xf32>, %arg1 : tensor<3x3x1x32xf32>) -> - tensor<100x28x28x1xf32> { - // expected-error@+1 {{expects padding-entries to have same dimension-size as size of window dimensions (2), but got: 3.}} +func.func @convolution_c14(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} %result = "stablehlo.convolution"(%arg0, %arg1) { batch_group_count = 1 : i64, dimension_numbers = #stablehlo.conv, %arg1 : ten >, feature_group_count = 1 : i64, lhs_dilation = array, - padding = dense<2> : tensor<3x2xi64>, + padding = dense<2> : tensor<2x2xi64>, rhs_dilation = array, window_strides = array } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> @@ -552,65 +501,67 @@ func.func @invalid_conv_dimensions(%arg0 : tensor<100x26x26x32xf32>, %arg1 : ten // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<0x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window to have positive value for 0-th window dimension, but got 0.}} +func.func @convolution_c15(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects input feature dimension (207) / feature_group_count = kernel input feature dimension (20). Got feature_group_count = 1.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<0x3x207x16xf32>) -> tensor<1x8x8x16xf32> + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<1x8x8x207xf32>, tensor<3x3x20x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window to have positive stride for 1-th window dimension, but got 0.}} +func.func @convolution_c16(%arg0: tensor<3x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<3x8x8x16xf32> { + // expected-error@+1 {{expects output feature dimension size (16) to be a multiple of batch_group_count. Got batch_group_count = 3.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 0], pad = [[1, 1], [1,1]], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { - batch_group_count = 1 : i64, + batch_group_count = 3 : i64, feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> - func.return %0 : tensor<1x8x8x16xf32> + precision_config = [#stablehlo, #stablehlo] + } : + (tensor<3x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<3x8x8x16xf32> + func.return %0 : tensor<3x8x8x16xf32> } // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window to have positive base dilation factor for 0-th window dimension, but got 0.}} +func.func @convolution_c17(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x69x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects kernel output feature dimension (16) to be divisible by feature_group_count. For feature_group_count = 3.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], - lhs_dilate = [0, 1], rhs_dilate = [1, 1]} + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, - feature_group_count = 1 : i64, + feature_group_count = 3 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x69x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- -func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, +func.func @convolution_c18(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+1 {{expects window to have positive window dilation factor for 0-th window dimension, but got 0.}} + // expected-error@+1 {{expects the same size for input, kernel and output spatial-dimensions, but got 2, 3, and 2 resp.}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], - lhs_dilate = [1, 1], rhs_dilate = [0, 1]} + dim_numbers = [b, 0, 1, f]x[0, 1, 2, i, o]->[b, 0, 1, f], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, @@ -622,225 +573,303 @@ func.func @invalid_conv_window_attributes(%arg0: tensor<1x8x8x207xf32>, // ----- -// Invalid rank of output-type. - -func.func @invalid_conv_return_type(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x16xf32> { - // expected-error @+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x8x16xf32>'}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1, 1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x16xf32> - func.return %0 : tensor<1x8x16xf32> +func.func @convolution_c19(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects kernel dimension-numbers to be unique, got {3, 2, 0, 0}.}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> } // ----- -// Invalid batch dimension in output-type. Should be equal to -// input-batch-dimension / batch_group_count. - -func.func @invalid_conv_return_type(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<2x8x8x16xf32> { - // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<2x8x8x16xf32>'}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<2x8x8x16xf32> - func.return %0 : tensor<2x8x8x16xf32> +func.func @convolution_c19(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> } // ----- -// Invalid feature dimension in output-type. Should be equal to -// kernel_output_feature_dimension. +func.func @convolution_c19(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> +} + +// ----- -func.func @invalid_conv_return_type(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x32xf32> { - // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x8x8x32xf32>'}} +func.func @convolution_c20(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects the same size for input, kernel and output spatial-dimensions, but got 2, 2, and 3 resp.}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], + dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, 2, f], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x32xf32> - func.return %0 : tensor<1x8x8x32xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + func.return %0 : tensor<1x8x8x16xf32> } // ----- -// The following tests checks the inferred output-type of ConvolutionOp. We -// deliberately put an invalid output-type in these tests so that the -// inffered-type can be highlighted in the error message. +func.func @convolution_c21(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects output dimension-numbers to be unique, got {0, 3, 0, 3}.}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> +} -// Dynamic input-batch-dimension -func.func @invalid_conv_dynamic_shapes(%arg0: tensor, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> { - // expected-error@+1 {{inferred shape '[?, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor, tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> - func.return %0 : tensor<1x1x1x1xf32> +// ----- + +func.func @convolution_c21(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> } // ----- -// Dynamic input-feature-dimension: No effect on output dimensions. -func.func @invalid_conv_dynamic_shapes(%arg0: tensor<1x8x8x?xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> { - // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x?xf32>, tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> - func.return %0 : tensor<1x1x1x1xf32> +func.func @convolution_c21(%arg0 : tensor<100x26x26x32xf32>, + %arg1 : tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> { + // expected-error@+1 {{expects input, kernel, and output dimension-numbers to be in-range [0, 4).}} + %result = "stablehlo.convolution"(%arg0, %arg1) { + batch_group_count = 1 : i64, + dimension_numbers = #stablehlo.conv, + feature_group_count = 1 : i64, + lhs_dilation = array, + padding = dense<2> : tensor<2x2xi64>, + rhs_dilation = array, + window_strides = array + } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>) -> + tensor<100x28x28x1xf32> + func.return %result : tensor<100x28x28x1xf32> } // ----- -// Dynamic input-spatial-dimension -func.func @invalid_conv_dynamic_shapes(%arg0: tensor<1x?x8x207xf32>, - %arg1: tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> { - // expected-error@+1 {{inferred shape '[1, ?, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} +func.func @convolution_c22(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects feature_group_count to be a positive number, got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { batch_group_count = 1 : i64, - feature_group_count = 1 : i64, + feature_group_count = 0 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x?x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x1x1x1xf32> - func.return %0 : tensor<1x1x1x1xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + func.return %0 : tensor<1x8x8x16xf32> } // ----- -// Dynamic kernel-input-feature-dimension: No effect on output dimensions. -func.func @invalid_conv_dynamic_shapes(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x?x16xf32>) -> tensor<1x1x1x1xf32> { - // expected-error@+1 {{inferred shape '[1, 8, 8, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} +func.func @convolution_c23(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects batch_group_count to be a positive number, got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { - batch_group_count = 1 : i64, + batch_group_count = 0 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x8x8x207xf32>, tensor<3x3x?x16xf32>) -> tensor<1x1x1x1xf32> - func.return %0 : tensor<1x1x1x1xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + func.return %0 : tensor<1x8x8x16xf32> } // ----- -// Dynamic kernel-output-feature-dimension -func.func @check_inferred_type_with_dynamic_input_dims(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207x?xf32>) -> tensor<1x1x1x1xf32> { - // expected-error@+1 {{inferred shape '[1, 8, 8, ?]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} +func.func @convolution_c24(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects batch_group_count and feature_group_count not to be both greater than 1. Got 2 and 2 resp.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, + batch_group_count = 2 : i64, + feature_group_count = 2 : i64, precision_config = [#stablehlo, #stablehlo] } : - (tensor<1x8x8x207xf32>, tensor<3x3x207x?xf32>) -> tensor<1x1x1x1xf32> - func.return %0 : tensor<1x1x1x1xf32> + (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + func.return %0 : tensor<1x8x8x16xf32> } // ----- -// Dynamic kernel-spatial-dimension -func.func @check_inferred_type_with_dynamic_input_dims(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x?x207x16xf32>) -> tensor<1x1x1x1xf32> { - // expected-error@+1 {{inferred shape '[1, 8, ?, 16]' is incompatible with return type of operation 'tensor<1x1x1x1xf32>'}} +func.func @convolution_c25(%arg0: tensor<3x2xf16>, + %arg1: tensor<2x2xf16>) -> tuple> { + // expected-error@+1{{expects precision config to be empty or have <= 2 elements}} %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1], [1,1]], - lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + dim_numbers = [b, f]x[i, o]->[b, f], + window = {stride = [], pad = [], lhs_dilate = [], rhs_dilate = [], + reverse = []} { batch_group_count = 1 : i64, feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo] - } : - (tensor<1x8x8x207xf32>, tensor<3x?x207x16xf32>) -> tensor<1x1x1x1xf32> - func.return %0 : tensor<1x1x1x1xf32> + precision_config = [#stablehlo, #stablehlo, #stablehlo] + } + : (tensor<3x2xf16>, tensor<2x2xf16>) -> tensor<3x2xf16> + %1 = "stablehlo.tuple"(%0) : (tensor<3x2xf16>) -> tuple> + func.return %1 : tuple> } // ----- -// This is an positive test in MLIR-HLO: -// https://github.com/tensorflow/mlir-hlo/blob/master/tests/Dialect/mhlo/ops.mlir#L3829 -// but negative here: stablehlo.convolution does no support unknown dimenstion -// dim_numbers = [b, 0, 1, ?, f]x[0, 1, ?, i, o]->[?, b, 0, 1, f] -// window = {stride = [1, 1], pad = [[1, 1], [1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} -func.func @conv2d_generic(%arg0: tensor<1x8x8x32x207xf32>, %arg1: tensor<3x3x32x207x16xf32>) -> tensor<32x1x8x8x16xf32> { - // expected-error@+1{{expects convolution arguments to have 4 dimensions. Got: 5}} - %0 = "stablehlo.convolution"(%arg0, %arg1) {batch_group_count = 1 : i64, - dimension_numbers = #stablehlo.conv, feature_group_count = 1 : i64, lhs_dilation = array, padding = dense<1> : tensor<2x2xi64>, precision_config = [#stablehlo, #stablehlo], rhs_dilation = array, window_strides = array} : - (tensor<1x8x8x32x207xf32>, tensor<3x3x32x207x16xf32>) -> tensor<32x1x8x8x16xf32> - func.return %0 : tensor<32x1x8x8x16xf32> +func.func @convolution_i4(%arg0: tensor<1x4x4x1xi64>, + %arg1: tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64> { + // expected-error@+1 {{expects the shape of padding-attribute to be {N, 2}, but got {2}.}} + %0 = "stablehlo.convolution"(%arg0, %arg1) { + window_strides = array, + padding = dense<0> : tensor<2xi64>, + lhs_dilation = array, + rhs_dilation = array, + window_reversal = array, + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, + feature_group_count = 1 : i64, + batch_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64> + func.return %0 : tensor<1x2x2x1xi64> } // ----- -func.func @conv2d(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { - // expected-error @+3 {{'stablehlo.convolution' Expected array with 2 elements, got 3 elements instead}} +func.func @convolution_invalid_window_attributes(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<0x3x207x16xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+1 {{expects window to have positive value for 0-th window dimension, but got 0.}} %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [1, 1], pad = [[1, 1, 1], [1, 1, 1]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} - {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : - (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> + window = {stride = [1, 1], pad = [[1, 1], [1,1]], + lhs_dilate = [1, 1], rhs_dilate = [1, 1]} + { + batch_group_count = 1 : i64, + feature_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo]} : + (tensor<1x8x8x207xf32>, tensor<0x3x207x16xf32>) -> tensor<1x8x8x16xf32> func.return %0 : tensor<1x8x8x16xf32> } // ----- // CHECK: module -// CHECK-SAME: stablehlo.conv = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 1, 0, f]> +// CHECK: stablehlo.conv = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 1, 0, f]> module attributes { stablehlo.conv = #stablehlo.conv[f, b, 0, 1]> } {} - -// ----- - -func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { - // expected-error@+3{{Unexpected keyword stide}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stide = [2, 1], pad = [[0, 1], [0, 1]], rhs_dilate = [1, 2]} - { batch_group_count = 1 : i64, feature_group_count = 1 : i64} - : (tensor<2x2x3x4xf32>, tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> - func.return %0 : tensor<3x5x5x4xf32> -} - -// ----- - -func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { - // expected-error@+3{{expected integer value}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [2, b], pad = [[0, 1], [0, 1]], rhs_dilate = [1, 2]} - { batch_group_count = 1 : i64, feature_group_count = 1 : i64} - : (tensor<2x2x3x4xf32>, tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> - func.return %0 : tensor<3x5x5x4xf32> -} - -// ----- - -func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { - // expected-error@+3{{Unexpected keyword stride}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], - window = {stride = [2, 1], pad = [[0, 1], [0, 1]], rhs_dilate = [1, 2], stride=[2,1]} - { batch_group_count = 1 : i64, feature_group_count = 1 : i64} - : (tensor<2x2x3x4xf32>, tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> - func.return %0 : tensor<3x5x5x4xf32> -} - -// ----- - -func.func @conv_invalid_precision_config(%arg0: tensor<3x2xf16>, - %arg1: tensor<2x2xf16>) -> tuple> { - // expected-error@+1{{expects precision config to be empty or have <= 2 elements}} - %0 = stablehlo.convolution(%arg0, %arg1) - dim_numbers = [b, f]x[i, o]->[b, f], - window = {stride = [], pad = [], lhs_dilate = [], rhs_dilate = [], - reverse = []} - { - batch_group_count = 1 : i64, - feature_group_count = 1 : i64, - precision_config = [#stablehlo, #stablehlo, #stablehlo] - } - : (tensor<3x2xf16>, tensor<2x2xf16>) -> tensor<3x2xf16> - %1 = "stablehlo.tuple"(%0) : (tensor<3x2xf16>) -> tuple> - func.return %1 : tuple> -}