Add interpreter for ConvolutionOp

openxla · Mar 11, 2023 · bc55ed7 · bc55ed7
1 parent 3722fce
commit bc55ed7
Show file tree

Hide file tree

Showing 62 changed files with 1,529 additions and 726 deletions.
diff --git a/docs/spec.md b/docs/spec.md
@@ -1943,41 +1943,42 @@ If `batch_group_count > 1`:
 #### Constraints
 
 <!-- markdownlint-disable line-length -->
-* (C1) $N =$ rank(`lhs`) $=$ rank(`rhs`).
-* (C2) element_type(`lhs`) $=$ element_type(`rhs`).
-* (C3) size(`window_strides`) $= N - 2$ .
-* (C4) `window_strides[i]` $\gt 0$  for all i $\in$ [0, size(`window_strides`)).
-* (C5) dim(`padding`, 0) $= N - 2$ and dim(`padding`, 1) = 2.
-* (C6) size(`lhs_dilation`) $= N - 2$.
-* (C7) `lhs_dilation[i]` $\gt 0$ for all i $\in$ [0, size(`lhs_dilation`)).
-* (C8) size(`rhs_dilation`) $= N - 2$.
-* (C9) `rhs_dilation[i]` $\gt 0$ for all i $\in$ [0, size(`rhs_dilation`)).
-* (C10) size(`window_reversal`) $= N - 2$.
-* (C11) `dim(lhs, input_batch_dimension) % batch_group_count = 0`.
-* (C12) `dim(lhs, input_feature_dimension) % feature_group_count = 0.
-* (C13) size(`input_spatial_dimensions`) $= N - 2$.
-* (C14) Given `input_dimensions = [input_batch_dimension] +
-       input_spatial_dimensions + [input_feature_dimension]`.
+* (C1) rank(`lhs`) $\ge$ 2.
+* (C2) $N =$ rank(`lhs`) $=$ rank(`rhs`).
+* (C3) element_type(`lhs`) $=$ element_type(`rhs`).
+* (C4) size(`window_strides`) $= N - 2$.
+* (C5) `window_strides[i]` $\gt 0$  for all i $\in$ [0, size(`window_strides`)).
+* (C6) dim(`padding`, 0) $= N - 2$ and dim(`padding`, 1) = 2.
+* (C7) size(`lhs_dilation`) $= N - 2$.
+* (C8) `lhs_dilation[i]` $\gt 0$ for all i $\in$ [0, size(`lhs_dilation`)).
+* (C9) size(`rhs_dilation`) $= N - 2$.
+* (C10) `rhs_dilation[i]` $\gt 0$ for all i $\in$ [0, size(`rhs_dilation`)).
+* (C11) size(`window_reversal`) $= N - 2$.
+* (C12) `dim(lhs, input_batch_dimension) % batch_group_count = 0`.
+* (C13) `dim(lhs, input_feature_dimension) % feature_group_count = 0`.
+* (C14) size(`input_spatial_dimensions`) $= N - 2$.
+* (C15) Given `input_dimensions = [input_batch_dimension] +
+       input_spatial_dimensions + [input_feature_dimension]`:
   * All dimensions in `input_dimensions` are unique.
   * For any i $\in$ `input_dimensions`, 0 $\le$ i $\lt$ N.
-* (C15) `dim(rhs, kernel_input_feature_dimension = dim(lhs, input_feature_dimension) / feature_group_count`.
-* (C16) `dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0`.
-* (C17) `dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0`.
-* (C18) size(`kernel_spatial_dimensions`) $= N - 2$.
-* (C19) Given `kernel_dimensions = kernel_spatial_dimensions +
-        [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`.
+* (C16) `dim(rhs, kernel_input_feature_dimension = dim(lhs, input_feature_dimension) / feature_group_count`.
+* (C17) `dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0`.
+* (C18) `dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0`.
+* (C19) size(`kernel_spatial_dimensions`) $= N - 2$.
+* (C20) Given `kernel_dimensions = kernel_spatial_dimensions +
+        [kernel_input_feature_dimension] + [kernel_output_feature_dimension]`:
   * All dimensions in `kernel_dimensions` are unique.
   * For any i $\in$ `kernel_dimensions`, 0 $\le$ i $\lt$ N.
-* (C20) size(`output_spatial_dimensions`) $= N - 2$.
-* (C21) Given `output_dimensions = [output_batch_dimension] +
-        output_spatial_dimensions + [output_feature_dimension]`.
+* (C21) size(`output_spatial_dimensions`) $= N - 2$.
+* (C22) Given `output_dimensions = [output_batch_dimension] +
+        output_spatial_dimensions + [output_feature_dimension]`:
   * All dimensions in `output_dimensions` are unique.
   * For any i $\in$ `output_dimensions`, 0 $\le$ i $\lt$ N.
-* (C22) `feature_group_count > 0`.
-* (C23) `batch_group_count > 0`.
-* (C24) `feature_group_count` $= 1$ OR  `batch_group_count` $= 1$.
-* (C25) size(`precision_config`) $=$ 2.
-* (C26) For result_dim $\in$ [0, N), `dim(result, result_dim)` is given by
+* (C23) `feature_group_count > 0`.
+* (C24) `batch_group_count > 0`.
+* (C25) Either `feature_group_count` $= 1$ or `batch_group_count` $= 1$.
+* (C26) size(`precision_config`) $=$ 2.
+* (C27) For result_dim $\in$ [0, N), `dim(result, result_dim)` is given by:
   * `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:
@@ -1988,8 +1989,6 @@ If `batch_group_count > 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`.
-* (C27) element_type(`result`) $=$ element_type(`lhs`).
-* (C28) rank(`result`) $= N$.
 <!-- markdownlint-enable line-length -->
 
 #### Examples
@@ -2030,13 +2029,15 @@ If `batch_group_count > 1`:
   feature_group_count = 1 : i64,
   batch_group_count = 1 : i64,
   precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
-} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>
+} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
 // %result: [[
 //            [[10], [26]],
 //            [[46], [62]]
 //          ]]
 ```
 
+&nbsp;[More Examples](../stablehlo/tests/interpret_convolution.mlir)
+
 ### cosine
 
 #### Semantics

diff --git a/docs/status.md b/docs/status.md
@@ -68,7 +68,7 @@ one of the following tracking labels.
 | concatenate              | yes           | yes          | yes            | yes             | yes         |
 | constant                 | yes           | yes          | yes            | yes             | yes         |
 | convert                  | yes           | yes          | infeasible     | yes             | no          |
-| convolution              | yes           | yes          | infeasible     | revisit         | no          |
+| convolution              | yes           | yes          | infeasible     | revisit         | yes         |
 | cosine                   | yes           | yes          | yes            | yes             | yes         |
 | count_leading_zeros      | yes           | yes          | yes            | yes             | no          |
 | create_token             | no            | yes\*        | yes\*          | yes             | no          |

diff --git a/output.mlir b/output.mlir
diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td
@@ -2037,24 +2037,32 @@ def StableHLO_ConvolutionOp : StableHLO_Op<"convolution", [Pure]> {
 
     Example:
     ```mlir
-    %result = "stablehlo.convolution"(%lhs, %rhs) {
-      window_strides = dense<4> : tensor<2xi64>,
-      padding = dense<0> : tensor<2x2xi64>,
-      lhs_dilation = dense<2> : tensor<2xi64>,
-      rhs_dilation = dense<1> : tensor<2xi64>,
-      window_reversal = dense<false> : tensor<2xi1>,
-      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<precision DEFAULT>, #stablehlo<precision DEFAULT>]
-    } : (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 = [false, false]
+      } {
+        feature_group_count = 1 : i64,
+        batch_group_count = 1 : i64,
+        precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
+      } :
+    (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
     ```
   }];
   let arguments = !con(
     (ins
-       HLO_Tensor:$lhs,
-       HLO_Tensor:$rhs),
-    StableHLO_ConvolutionAttributes.attributes);
+       HLO_Tensor:$lhs, /*convolution_i1*/
+       HLO_Tensor:$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_Tensor);
   let hasVerifier = 1;