Index: mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml =================================================================== --- mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml +++ mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml @@ -1643,6 +1643,107 @@ operands: - !ScalarExpression scalar_arg: K + +--- !LinalgOpConfig +metadata: !LinalgOpMetadata + name: conv_2d_gnchw_gfchw + cpp_class_name: Conv2DNgchwFgchwOp + doc: |- + Performs 2-D convolution. + + Layout: + * Input: NGCHW. + * Kernel: FGCHW. + + 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] -> (s0, + s1, s2, s3 * s4 + s5 * s6, s7 * s8 + s9 * s10)> + - !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] -> (s11, + s1, s5, s9)> + - !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] -> (s0, + s11, s3, s7)> + - !LinalgOperandDefConfig + name: strides + kind: index_attr + index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] -> + (s4, s8)> + default_indices: + - 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] -> + (s6, s10)> + default_indices: + - 1 + - 1 + indexing_maps: !LinalgIndexingMapsConfig + static_indexing_maps: + - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8, + s9, s10, s11] -> (d0, d2, d5, d3 * s4 + d6 * s6, d4 * s8 + d7 * s10)> + - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8, + s9, s10, s11] -> (d1, d2, d5, d6, d7)> + - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8, + s9, s10, s11] -> (d0, d1, d3, d4)> + iterator_types: + - 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: conv_3d_ndhwc_dhwcf Index: mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py =================================================================== --- mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py +++ mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py @@ -298,7 +298,7 @@ S.OW * S.SW + S.KW * S.DW, S.C), K=TensorDef(T2, S.KH, S.KW, S.C, S.F), IZp=ScalarDef(I32), - KZp=ScalarDef(I32), + KgZp=ScalarDef(I32), O=TensorDef(U, S.N, S.OH, S.OW, S.F, output=True), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): @@ -343,6 +343,28 @@ U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.f, D.c, D.kh, D.kw]) +@linalg_structured_op +def conv_group_2d_ngchw_fgchw(I=TensorDef(T1, S.N, G, S.CG, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW), + K=TensorDef(T2, S.FG, G, S.C, S.KH, S.KW), + O=TensorDef(U, S.N, S.F, S.OH, S.OW, output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): + """Performs Grouped 2-D convolution. + + Layout: + * Input: NGCHW. + * Kernel: FGCHW. + + 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.fg, D.oh, D.ow, D.cg, D.c, D.kh, D.kw, D.g) + O[D.n, D.fg, D.oh, D.ow] += TypeFn.cast_signed( + U, I[D.n, D.g, D.cg, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + + D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.fg, D.g, D.c, D.kh, D.kw]) + @linalg_structured_op def conv_3d_ndhwc_dhwcf(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD,