diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml @@ -1513,6 +1513,97 @@ - !ScalarExpression scalar_arg: K --- !LinalgOpConfig +metadata: !LinalgOpMetadata + name: conv_1d_ngcw_gfcw + cpp_class_name: Conv1DNgcwGfcwOp + doc: |- + Performs 1-D grouped convolution. + + Layout: + * Input: NGCW. + * Kernel: GFCW. + + Numeric casting is performed on the operands to the inner multiply, promoting + them to the same data type as the accumulator/output. + implements: + - LinalgConvolutionOpInterface +structured_op: !LinalgStructuredOpConfig + args: + - !LinalgOperandDefConfig + name: I + kind: input_tensor + type_var: T1 + shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7] -> (s0, s1, s2, s3 * + s4 + s5 * s6)> + - !LinalgOperandDefConfig + name: K + kind: input_tensor + type_var: T2 + shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7] -> (s1, s7, s2, s5)> + - !LinalgOperandDefConfig + name: O + kind: output_tensor + type_var: U + shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7] -> (s0, s1, s7, s3)> + - !LinalgOperandDefConfig + name: strides + kind: index_attr + index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7] -> (s4)> + default_indices: + - 1 + - !LinalgOperandDefConfig + name: dilations + kind: index_attr + index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7] -> (s6)> + default_indices: + - 1 + indexing_maps: !LinalgIndexingMapsConfig + static_indexing_maps: + - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7] -> (d0, + d1, d4, d3 * s4 + d5 * s6)> + - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7] -> (d1, + d2, d4, d5)> + - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7] -> (d0, + d1, d2, d3)> + iterator_types: + - parallel + - parallel + - parallel + - parallel + - reduction + - reduction + assignments: + - !ScalarAssign + arg: O + value: !ScalarExpression + scalar_fn: + kind: binary + fn_name: add + operands: + - !ScalarExpression + scalar_arg: O + - !ScalarExpression + scalar_fn: + kind: binary + fn_name: mul + operands: + - !ScalarExpression + scalar_fn: + kind: type + fn_name: cast_signed + type_var: U + operands: + - !ScalarExpression + scalar_arg: I + - !ScalarExpression + scalar_fn: + kind: type + fn_name: cast_signed + type_var: U + operands: + - !ScalarExpression + scalar_arg: K +--- !LinalgOpConfig metadata: !LinalgOpMetadata name: conv_2d_nhwc_hwcf cpp_class_name: Conv2DNhwcHwcfOp @@ -2282,6 +2373,111 @@ - !ScalarExpression scalar_arg: KZp --- !LinalgOpConfig +metadata: !LinalgOpMetadata + name: conv_3d_ngcdhw_gfcdhw + cpp_class_name: Conv3DNgcdhwGfcdhwOp + doc: |- + Performs 3-D grouped convolution. + + Layout: + * Input: NGCDHW. + * Kernel: GFCDHW. + + Numeric casting is performed on the operands to the inner multiply, promoting + them to the same data type as the accumulator/output. + implements: + - LinalgConvolutionOpInterface +structured_op: !LinalgStructuredOpConfig + args: + - !LinalgOperandDefConfig + name: I + kind: input_tensor + type_var: T1 + shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, + s13, s14, s15] -> (s0, s1, s2, s3 * s4 + s5 * s6, s7 * s8 + s9 * s10, s11 * + s12 + s13 * s14)> + - !LinalgOperandDefConfig + name: K + kind: input_tensor + type_var: T2 + shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, + s13, s14, s15] -> (s1, s15, s2, s5, s9, s13)> + - !LinalgOperandDefConfig + name: O + kind: output_tensor + type_var: U + shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, + s13, s14, s15] -> (s0, s1, s15, s3, s7, s11)> + - !LinalgOperandDefConfig + name: strides + kind: index_attr + index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, + s12, s13, s14, s15] -> (s4, s8, s12)> + default_indices: + - 1 + - 1 + - 1 + - !LinalgOperandDefConfig + name: dilations + kind: index_attr + index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, + s12, s13, s14, s15] -> (s6, s10, s14)> + default_indices: + - 1 + - 1 + - 1 + indexing_maps: !LinalgIndexingMapsConfig + static_indexing_maps: + - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9)[s0, s1, s2, s3, s4, s5, + s6, s7, s8, s9, s10, s11, s12, s13, s14, s15] -> (d0, d1, d6, d3 * s4 + d7 * + s6, d4 * s8 + d8 * s10, d5 * s12 + d9 * s14)> + - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9)[s0, s1, s2, s3, s4, s5, + s6, s7, s8, s9, s10, s11, s12, s13, s14, s15] -> (d1, d2, d6, d7, d8, d9)> + - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9)[s0, s1, s2, s3, s4, s5, + s6, s7, s8, s9, s10, s11, s12, s13, s14, s15] -> (d0, d1, d2, d3, d4, d5)> + iterator_types: + - parallel + - parallel + - parallel + - parallel + - parallel + - parallel + - reduction + - reduction + - reduction + - reduction + assignments: + - !ScalarAssign + arg: O + value: !ScalarExpression + scalar_fn: + kind: binary + fn_name: add + operands: + - !ScalarExpression + scalar_arg: O + - !ScalarExpression + scalar_fn: + kind: binary + fn_name: mul + operands: + - !ScalarExpression + scalar_fn: + kind: type + fn_name: cast_signed + type_var: U + operands: + - !ScalarExpression + scalar_arg: I + - !ScalarExpression + scalar_fn: + kind: type + fn_name: cast_signed + type_var: U + operands: + - !ScalarExpression + scalar_arg: K +--- !LinalgOpConfig metadata: !LinalgOpMetadata name: depthwise_conv_1d_nwc_wc cpp_class_name: DepthwiseConv1DNwcWcOp diff --git a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py --- a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py +++ b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py @@ -308,6 +308,28 @@ U, K[D.f, D.c, D.kw]) +@linalg_structured_op +def conv_1d_ngcw_gfcw(I=TensorDef(T1, S.N, S.G, S.C, S.OW * S.SW + S.KW * S.DW), + K=TensorDef(T2, S.G, S.FG, S.C, S.KW), + O=TensorDef(U, S.N, S.G, S.FG, S.OW, output=True), + strides=IndexAttrDef(S.SW, default=[1]), + dilations=IndexAttrDef(S.DW, default=[1])): + """Performs 1-D grouped convolution. + + Layout: + * Input: NGCW. + * Kernel: GFCW. + + Numeric casting is performed on the operands to the inner multiply, promoting + them to the same data type as the accumulator/output. + """ + implements(ConvolutionOpInterface) + domain(D.n, D.g, D.fg, D.ow, D.c, D.kw) + O[D.n, D.g, D.fg, D.ow] += TypeFn.cast_signed( + U, I[D.n, D.g, D.c, D.ow * S.SW + D.kw * S.DW]) * TypeFn.cast_signed( + U, K[D.g, D.fg, D.c, D.kw]) + + @linalg_structured_op def conv_2d_nhwc_hwcf(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C), @@ -493,6 +515,45 @@ TypeFn.cast_signed(U, KZp)) +@linalg_structured_op +def conv_3d_ngcdhw_gfcdhw(I=TensorDef(T1, S.N, S.G, S.C, + S.OD * S.SD + S.KD * S.DD, + S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW), + K=TensorDef(T2, S.G, S.FG, S.C, S.KD, S.KH, S.KW), + O=TensorDef(U, + S.N, + S.G, + S.FG, + S.OD, + S.OH, + S.OW, + output=True), + strides=IndexAttrDef(S.SD, + S.SH, + S.SW, + default=[1, 1, 1]), + dilations=IndexAttrDef(S.DD, + S.DH, + S.DW, + default=[1, 1, 1])): + """Performs 3-D grouped convolution. + + Layout: + * Input: NGCDHW. + * Kernel: GFCDHW. + + Numeric casting is performed on the operands to the inner multiply, promoting + them to the same data type as the accumulator/output. + """ + implements(ConvolutionOpInterface) + domain(D.n, D.g, D.fg, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw) + O[D.n, D.g, D.fg, D.od, D.oh, D.ow] += TypeFn.cast_signed( + U, I[D.n, D.g, D.c, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, + D.ow * S.SW + D.kw * S.DW]) * TypeFn.cast_signed( + U, K[D.g, D.fg, D.c, D.kd, D.kh, D.kw]) + + @linalg_structured_op def depthwise_conv_1d_nwc_wc(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.IC), diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir --- a/mlir/test/Dialect/Linalg/named-ops.mlir +++ b/mlir/test/Dialect/Linalg/named-ops.mlir @@ -285,6 +285,54 @@ // ----- +// CHECK-LABEL: func @conv_1d_ngcw_gfcw +func.func @conv_1d_ngcw_gfcw(%input: tensor, %filter: tensor, %init: tensor) -> tensor { + // CHECK: %{{.+}} = linalg.conv_1d_ngcw_gfcw + // CHECK-SAME: dilations = dense<1> : tensor<1xi64> + // CHECK-SAME: strides = dense<1> : tensor<1xi64> + // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor, tensor) + // CHECK-SAME: outs(%{{.+}} : tensor) -> tensor + %0 = linalg.conv_1d_ngcw_gfcw {dilations = dense<1> : tensor<1xi64>, + strides = dense<1> : tensor<1xi64>} + ins (%input, %filter: tensor, tensor) + outs (%init: tensor) -> tensor + return %0 : tensor +} + +// ----- + +// CHECK-LABEL: func @conv_1d_ngcw_gfcw_static +func.func @conv_1d_ngcw_gfcw_static(%input: tensor<1x2x32x128xf32>, %filter: tensor<2x64x32x3xf32>, %init: tensor<1x2x64x126xf32>) -> tensor<1x2x64x126xf32> { + // CHECK: %{{.+}} = linalg.conv_1d_ngcw_gfcw + // CHECK-SAME: dilations = dense<1> : tensor<1xi64> + // CHECK-SAME: strides = dense<1> : tensor<1xi64> + // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x2x32x128xf32>, tensor<2x64x32x3xf32>) + // CHECK-SAME: outs(%{{.+}} : tensor<1x2x64x126xf32>) -> tensor<1x2x64x126xf32> + %0 = linalg.conv_1d_ngcw_gfcw {dilations = dense<1> : tensor<1xi64>, + strides = dense<1> : tensor<1xi64>} + ins (%input, %filter: tensor<1x2x32x128xf32>, tensor<2x64x32x3xf32>) + outs (%init: tensor<1x2x64x126xf32>) -> tensor<1x2x64x126xf32> + return %0 : tensor<1x2x64x126xf32> +} + +// ----- + +// CHECK-LABEL: func @conv_1d_ngcw_gfcw +func.func @conv_1d_ngcw_gfcw(%input: memref, %filter: memref, %output: memref) { + // CHECK: linalg.conv_1d_ngcw_gfcw + // CHECK-SAME: dilations = dense<1> : tensor<1xi64> + // CHECK-SAME: strides = dense<1> : tensor<1xi64> + // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref, memref) + // CHECK-SAME: outs(%{{.+}} : memref) + linalg.conv_1d_ngcw_gfcw {dilations = dense<1> : tensor<1xi64>, + strides = dense<1> : tensor<1xi64>} + ins (%input, %filter: memref, memref) + outs (%output: memref) + return +} + +// ----- + // CHECK-LABEL: func @conv_2d_nhwc_hwcf func.func @conv_2d_nhwc_hwcf(%input: tensor, %filter: tensor, %init: tensor) -> tensor { // CHECK: %{{.+}} = linalg.conv_2d_nhwc_hwcf @@ -429,6 +477,54 @@ // ----- +// CHECK-LABEL: func @conv_3d_ngcdhw_gfcdhw +func.func @conv_3d_ngcdhw_gfcdhw(%input: tensor, %filter: tensor, %init: tensor) -> tensor { + // CHECK: %{{.+}} = linalg.conv_3d_ngcdhw_gfcdhw + // CHECK-SAME: dilations = dense<1> : tensor<3xi64> + // CHECK-SAME: strides = dense<1> : tensor<3xi64> + // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor, tensor) + // CHECK-SAME: outs(%{{.+}} : tensor) -> tensor + %0 = linalg.conv_3d_ngcdhw_gfcdhw {dilations = dense<1> : tensor<3xi64>, + strides = dense<1> : tensor<3xi64>} + ins (%input, %filter: tensor, tensor) + outs (%init: tensor) -> tensor + return %0 : tensor +} + +// ----- + +// CHECK-LABEL: func @conv_3d_ngcdhw_gfcdhw_static +func.func @conv_3d_ngcdhw_gfcdhw_static(%input: tensor<1x2x32x128x256x512xf32>, %filter: tensor<2x64x32x3x5x7xf32>, %init: tensor<1x2x64x126x252x506xf32>) -> tensor<1x2x64x126x252x506xf32> { + // CHECK: %{{.+}} = linalg.conv_3d_ngcdhw_gfcdhw + // CHECK-SAME: dilations = dense<1> : tensor<3xi64> + // CHECK-SAME: strides = dense<1> : tensor<3xi64> + // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x2x32x128x256x512xf32>, tensor<2x64x32x3x5x7xf32>) + // CHECK-SAME: outs(%{{.+}} : tensor<1x2x64x126x252x506xf32>) -> tensor<1x2x64x126x252x506xf32> + %0 = linalg.conv_3d_ngcdhw_gfcdhw {dilations = dense<1> : tensor<3xi64>, + strides = dense<1> : tensor<3xi64>} + ins (%input, %filter: tensor<1x2x32x128x256x512xf32>, tensor<2x64x32x3x5x7xf32>) + outs (%init: tensor<1x2x64x126x252x506xf32>) -> tensor<1x2x64x126x252x506xf32> + return %0 : tensor<1x2x64x126x252x506xf32> +} + +// ----- + +// CHECK-LABEL: func @conv_3d_ngcdhw_gfcdhw +func.func @conv_3d_ngcdhw_gfcdhw(%input: memref, %filter: memref, %output: memref) { + // CHECK: linalg.conv_3d_ngcdhw_gfcdhw + // CHECK-SAME: dilations = dense<1> : tensor<3xi64> + // CHECK-SAME: strides = dense<1> : tensor<3xi64> + // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref, memref) + // CHECK-SAME: outs(%{{.+}} : memref) + linalg.conv_3d_ngcdhw_gfcdhw {dilations = dense<1> : tensor<3xi64>, + strides = dense<1> : tensor<3xi64>} + ins (%input, %filter: memref, memref) + outs (%output: memref) + return +} + +// ----- + // CHECK-LABEL: func @pooling_nhwc_sum_tensor // CHECK: %{{.+}} = linalg.pooling_nhwc_sum // CHECK-SAME: dilations = dense<1> : tensor<2xi64>