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 @@ -1,6 +1,47 @@ ### AUTOGENERATED from core_named_ops.py ### To regenerate, run: bin/update_core_linalg_named_ops.sh --- !LinalgOpConfig +metadata: !LinalgOpMetadata + name: copy + cpp_class_name: CopyOp + doc: |- + Copies the tensor elementwise. + + Numeric casting is performed on the input operand, promoting it to the same + data type as the accumulator/output. +structured_op: !LinalgStructuredOpConfig + args: + - !LinalgOperandDefConfig + name: I + kind: input_tensor + type_var: T1 + shape_map: affine_map<() -> ()> + - !LinalgOperandDefConfig + name: O + kind: output_tensor + type_var: U + shape_map: affine_map<() -> ()> + - !LinalgOperandDefConfig + name: cast + kind: type_fn_attr + default_fn: cast_signed + indexing_maps: !LinalgIndexingMapsConfig + static_indexing_maps: + - affine_map<() -> ()> + - affine_map<() -> ()> + iterator_types: [] + assignments: + - !ScalarAssign + arg: O + value: !ScalarExpression + scalar_fn: + kind: type + attr_name: cast + type_var: U + operands: + - !ScalarExpression + scalar_arg: I +--- !LinalgOpConfig metadata: !LinalgOpMetadata name: elemwise_unary cpp_class_name: ElemwiseUnaryOp 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 @@ -7,11 +7,22 @@ @linalg_structured_op -def elemwise_unary( - I=TensorDef(T1), - O=TensorDef(U, output=True), - fun=UnaryFnAttrDef(default=UnaryFn.exp), - cast=TypeFnAttrDef(default=TypeFn.cast_signed)): +def copy(I=TensorDef(T1), + O=TensorDef(U, output=True), + cast=TypeFnAttrDef(default=TypeFn.cast_signed)): + """Copies the tensor elementwise. + + Numeric casting is performed on the input operand, promoting it to the same + data type as the accumulator/output. + """ + O[None] = cast(U, I[None]) + + +@linalg_structured_op +def elemwise_unary(I=TensorDef(T1), + O=TensorDef(U, output=True), + fun=UnaryFnAttrDef(default=UnaryFn.exp), + cast=TypeFnAttrDef(default=TypeFn.cast_signed)): """Applies the unary function fun elementwise. Numeric casting is performed on the input operand, promoting it to the same @@ -21,12 +32,11 @@ @linalg_structured_op -def elemwise_binary( - lhs=TensorDef(T1), - rhs=TensorDef(T2), - O=TensorDef(U, output=True), - fun=BinaryFnAttrDef(default=BinaryFn.add), - cast=TypeFnAttrDef(default=TypeFn.cast_signed)): +def elemwise_binary(lhs=TensorDef(T1), + rhs=TensorDef(T2), + O=TensorDef(U, output=True), + fun=BinaryFnAttrDef(default=BinaryFn.add), + cast=TypeFnAttrDef(default=TypeFn.cast_signed)): """Applies the binary function fun elementwise. Numeric casting is performed on the input operand, promoting it to the same @@ -36,11 +46,10 @@ @linalg_structured_op -def matmul( - A=TensorDef(T1, S.M, S.K), - B=TensorDef(T2, S.K, S.N), - C=TensorDef(U, S.M, S.N, output=True), - cast=TypeFnAttrDef(default=TypeFn.cast_signed)): +def matmul(A=TensorDef(T1, S.M, S.K), + B=TensorDef(T2, S.K, S.N), + C=TensorDef(U, S.M, S.N, output=True), + cast=TypeFnAttrDef(default=TypeFn.cast_signed)): """Performs a matrix multiplication of two 2D inputs. Numeric casting is performed on the operands to the inner multiply, promoting @@ -52,10 +61,9 @@ @linalg_structured_op -def matmul_unsigned( - A=TensorDef(T1, S.M, S.K), - B=TensorDef(T2, S.K, S.N), - C=TensorDef(U, S.M, S.N, output=True)): +def matmul_unsigned(A=TensorDef(T1, S.M, S.K), + B=TensorDef(T2, S.K, S.N), + C=TensorDef(U, S.M, S.N, output=True)): """Performs an unsigned matrix multiplication of two 2D inputs. Numeric casting is performed on the operands to the inner multiply, promoting @@ -68,12 +76,11 @@ @linalg_structured_op -def quantized_matmul( - A=TensorDef(T1, S.M, S.K), - B=TensorDef(T2, S.K, S.N), - AZp=ScalarDef(I32), - BZp=ScalarDef(I32), - C=TensorDef(U, S.M, S.N, output=True)): +def quantized_matmul(A=TensorDef(T1, S.M, S.K), + B=TensorDef(T2, S.K, S.N), + AZp=ScalarDef(I32), + BZp=ScalarDef(I32), + C=TensorDef(U, S.M, S.N, output=True)): """Performs a matrix multiplication of two 2D inputs. Numeric casting is performed on the operands to the inner multiply, promoting @@ -82,16 +89,16 @@ matmul. """ domain(D.m, D.n, D.k) - C[D.m, D.n] += ( - TypeFn.cast_signed(U, A[D.m, D.k]) - TypeFn.cast_signed(U, AZp)) * ( - TypeFn.cast_signed(U, B[D.k, D.n]) - TypeFn.cast_signed(U, BZp)) + C[D.m, + D.n] += (TypeFn.cast_signed(U, A[D.m, D.k]) - + TypeFn.cast_signed(U, AZp)) * (TypeFn.cast_signed(U, B[D.k, D.n]) - + TypeFn.cast_signed(U, BZp)) @linalg_structured_op -def mmt4d( - lhs=TensorDef(TV.LhsType, S.M, S.K, S.M0, S.K0), - rhs=TensorDef(TV.RhsType, S.N, S.K, S.N0, S.K0), - accum=TensorDef(TV.AccumType, S.M, S.N, S.M0, S.N0, output=True)): +def mmt4d(lhs=TensorDef(TV.LhsType, S.M, S.K, S.M0, S.K0), + rhs=TensorDef(TV.RhsType, S.N, S.K, S.N0, S.K0), + accum=TensorDef(TV.AccumType, S.M, S.N, S.M0, S.N0, output=True)): """Performs a matrix-matrix-transpose multiplication of two 4D inputs. Differences from linalg.matmul: @@ -110,10 +117,9 @@ @linalg_structured_op -def batch_matmul( - A=TensorDef(T1, Batch, S.M, S.K), - B=TensorDef(T2, Batch, S.K, S.N), - C=TensorDef(U, Batch, S.M, S.N, output=True)): +def batch_matmul(A=TensorDef(T1, Batch, S.M, S.K), + B=TensorDef(T2, Batch, S.K, S.N), + C=TensorDef(U, Batch, S.M, S.N, output=True)): """Performs a batched matrix multiplication of two 3D inputs. Numeric casting is performed on the operands to the inner multiply, promoting @@ -127,12 +133,11 @@ @linalg_structured_op -def quantized_batch_matmul( - A=TensorDef(T1, Batch, S.M, S.K), - B=TensorDef(T2, Batch, S.K, S.N), - AZp=ScalarDef(I32), - BZp=ScalarDef(I32), - C=TensorDef(U, Batch, S.M, S.N, output=True)): +def quantized_batch_matmul(A=TensorDef(T1, Batch, S.M, S.K), + B=TensorDef(T2, Batch, S.K, S.N), + AZp=ScalarDef(I32), + BZp=ScalarDef(I32), + C=TensorDef(U, Batch, S.M, S.N, output=True)): """Performs a batched matrix multiplication of two 3D inputs. Numeric casting is performed on the operands to the inner multiply, promoting @@ -141,16 +146,15 @@ matmul. """ domain(D.b, D.m, D.n, D.k) - C[D.b, D.m, D.n] += ( - TypeFn.cast_signed(U, A[D.b, D.m, D.k]) - TypeFn.cast_signed(U, AZp)) * ( - TypeFn.cast_signed(U, B[D.b, D.k, D.n]) - TypeFn.cast_signed(U, BZp)) + C[D.b, D.m, D.n] += (TypeFn.cast_signed(U, A[D.b, D.m, D.k]) - + TypeFn.cast_signed(U, AZp)) * (TypeFn.cast_signed( + U, B[D.b, D.k, D.n]) - TypeFn.cast_signed(U, BZp)) @linalg_structured_op -def matvec( - A=TensorDef(T1, S.M, S.N), - y=TensorDef(T2, S.N), - x=TensorDef(U, S.M, output=True)): +def matvec(A=TensorDef(T1, S.M, S.N), + y=TensorDef(T2, S.N), + x=TensorDef(U, S.M, output=True)): """Performs a matrix-vector multiplication. Numeric casting is performed on the operands to the inner multiply, promoting @@ -162,10 +166,9 @@ @linalg_structured_op -def vecmat( - y=TensorDef(T1, S.M), - A=TensorDef(T2, S.M, S.N), - x=TensorDef(U, S.N, output=True)): +def vecmat(y=TensorDef(T1, S.M), + A=TensorDef(T2, S.M, S.N), + x=TensorDef(U, S.N, output=True)): """Performs a vector-matrix multiplication. Numeric casting is performed on the operands to the inner multiply, promoting @@ -177,10 +180,9 @@ @linalg_structured_op -def batch_matvec( - A=TensorDef(T1, Batch, S.M, S.K), - B=TensorDef(T2, Batch, S.K), - C=TensorDef(U, Batch, S.M, output=True)): +def batch_matvec(A=TensorDef(T1, Batch, S.M, S.K), + B=TensorDef(T2, Batch, S.K), + C=TensorDef(U, Batch, S.M, output=True)): """Performs a batched matrix-vector multiplication. Numeric casting is performed on the operands to the inner multiply, promoting @@ -193,8 +195,8 @@ @linalg_structured_op -def dot( - A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U, output=True)): +def dot(A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U, + output=True)): """Performs a dot product of two vectors to a scalar result. Numeric casting is performed on the operands to the inner multiply, promoting @@ -205,10 +207,9 @@ @linalg_structured_op -def conv_1d( - I=TensorDef(T1, S.OW + S.KW), - K=TensorDef(T2, S.KW), - O=TensorDef(U, S.OW, output=True)): +def conv_1d(I=TensorDef(T1, S.OW + S.KW), + K=TensorDef(T2, S.KW), + O=TensorDef(U, S.OW, output=True)): """Performs 1-D convolution with no channels. Numeric casting is performed on the operands to the inner multiply, promoting @@ -221,10 +222,9 @@ @linalg_structured_op -def conv_2d( - I=TensorDef(T1, S.OH + S.KH, S.OW + S.KW), - K=TensorDef(T2, S.KH, S.KW), - O=TensorDef(U, S.OH, S.OW, output=True)): +def conv_2d(I=TensorDef(T1, S.OH + S.KH, S.OW + S.KW), + K=TensorDef(T2, S.KH, S.KW), + O=TensorDef(U, S.OH, S.OW, output=True)): """Performs 2-D convolution with no channels. Numeric casting is performed on the operands to the inner multiply, promoting @@ -237,10 +237,9 @@ @linalg_structured_op -def conv_3d( - I=TensorDef(T1, S.OD + S.KD, S.OH + S.KH, S.OW + S.KW), - K=TensorDef(T2, S.KD, S.KH, S.KW), - O=TensorDef(U, S.OD, S.OH, S.OW, output=True)): +def conv_3d(I=TensorDef(T1, S.OD + S.KD, S.OH + S.KH, S.OW + S.KW), + K=TensorDef(T2, S.KD, S.KH, S.KW), + O=TensorDef(U, S.OD, S.OH, S.OW, output=True)): """Performs 3-D convolution with no channels. Numeric casting is performed on the operands to the inner multiply, promoting @@ -254,12 +253,11 @@ @linalg_structured_op -def conv_1d_nwc_wcf( - I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C), - K=TensorDef(T2, S.KW, S.C, S.F), - O=TensorDef(U, S.N, S.OW, S.F, output=True), - strides=IndexAttrDef(S.SW, default=[1]), - dilations=IndexAttrDef(S.DW, default=[1])): +def conv_1d_nwc_wcf(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, S.KW, S.C, S.F), + O=TensorDef(U, S.N, S.OW, S.F, output=True), + strides=IndexAttrDef(S.SW, default=[1]), + dilations=IndexAttrDef(S.DW, default=[1])): """Performs 1-D convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -267,19 +265,18 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.ow, D.f, D.kw, D.c) - O[D.n, D.ow, - D.f] += TypeFn.cast_signed(U, I[D.n, D.ow * S.SW + D.kw * S.DW, - D.c]) * TypeFn.cast_signed(U, K[D.kw, D.c, D.f]) + O[D.n, D.ow, D.f] += TypeFn.cast_signed( + U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]) * TypeFn.cast_signed( + U, K[D.kw, D.c, D.f]) @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), - K=TensorDef(T2, S.KH, S.KW, S.C, S.F), - 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])): +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), + K=TensorDef(T2, S.KH, S.KW, S.C, S.F), + 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])): """Performs 2-D convolution. Layout: @@ -297,15 +294,14 @@ @linalg_structured_op -def conv_2d_nhwc_hwcf_q( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, 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), - 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])): +def conv_2d_nhwc_hwcf_q(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + 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), + 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])): """Performs 2-D convolution with zero point offsets. Layout: @@ -321,19 +317,17 @@ O[D.n, D.oh, D.ow, D.f] += (TypeFn.cast_signed( U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]) - - TypeFn.cast_signed(U, IZp)) * ( - TypeFn.cast_signed(U, K[D.kh, D.kw, D.c, D.f]) - - TypeFn.cast_signed(U, KZp)) + TypeFn.cast_signed(U, IZp)) * (TypeFn.cast_signed( + U, K[D.kh, D.kw, D.c, D.f]) - TypeFn.cast_signed(U, KZp)) @linalg_structured_op -def conv_2d_nchw_fchw( - I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, - S.OW * S.SW + S.KW * S.DW), - K=TensorDef(T2, S.F, 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])): +def conv_2d_nchw_fchw(I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW), + K=TensorDef(T2, S.F, 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 2-D convolution. Layout: @@ -351,13 +345,19 @@ @linalg_structured_op -def conv_3d_ndhwc_dhwcf( - I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, - S.OW * S.SW + S.KW * S.DW, S.C), - K=TensorDef(T2, S.KD, S.KH, S.KW, S.C, S.F), - O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.F, 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])): +def conv_3d_ndhwc_dhwcf(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, + S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, S.KD, S.KH, S.KW, S.C, S.F), + O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.F, 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 convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -372,12 +372,12 @@ @linalg_structured_op -def depthwise_conv_1d_nwc_wc( - I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.IC), - K=TensorDef(T2, S.KW, S.IC), - O=TensorDef(U, S.N, S.OW, S.IC, output=True), - strides=IndexAttrDef(S.SW, default=[1]), - dilations=IndexAttrDef(S.DW, default=[1])): +def depthwise_conv_1d_nwc_wc(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, + S.IC), + K=TensorDef(T2, S.KW, S.IC), + O=TensorDef(U, S.N, S.OW, S.IC, output=True), + strides=IndexAttrDef(S.SW, default=[1]), + dilations=IndexAttrDef(S.DW, default=[1])): """Performs depth-wise 1-D convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -392,13 +392,19 @@ @linalg_structured_op -def depthwise_conv_2d_nhwc_hwc( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.IC), - K=TensorDef(T2, S.KH, S.KW, S.IC), - O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def depthwise_conv_2d_nhwc_hwc(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.IC), + K=TensorDef(T2, S.KH, S.KW, S.IC), + O=TensorDef(U, + S.N, + S.OH, + S.OW, + S.IC, + output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, + S.DW, + default=[1, 1])): """Performs depth-wise 2-D convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -413,15 +419,23 @@ @linalg_structured_op -def depthwise_conv_2d_nhwc_hwc_q( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.IC), - K=TensorDef(T2, S.KH, S.KW, S.IC), - IZp=ScalarDef(I32), - KZp=ScalarDef(I32), - O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def depthwise_conv_2d_nhwc_hwc_q(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.IC), + K=TensorDef(T2, S.KH, S.KW, S.IC), + IZp=ScalarDef(I32), + KZp=ScalarDef(I32), + O=TensorDef(U, + S.N, + S.OH, + S.OW, + S.IC, + output=True), + strides=IndexAttrDef(S.SH, + S.SW, + default=[1, 1]), + dilations=IndexAttrDef(S.DH, + S.DW, + default=[1, 1])): """Performs depth-wise 2-D convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -437,13 +451,21 @@ @linalg_structured_op -def depthwise_conv_2d_nhwc_hwcm( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.IC), - K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM), - O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def depthwise_conv_2d_nhwc_hwcm(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.IC), + K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM), + O=TensorDef(U, + S.N, + S.OH, + S.OW, + S.IC, + S.CM, + output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, + 1]), + dilations=IndexAttrDef(S.DH, + S.DW, + default=[1, 1])): """Performs depth-wise 2-D convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -457,15 +479,25 @@ @linalg_structured_op -def depthwise_conv_2d_nhwc_hwcm_q( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.IC), - K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM), - IZp=ScalarDef(I32), - KZp=ScalarDef(I32), - O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def depthwise_conv_2d_nhwc_hwcm_q(I=TensorDef(T1, S.N, + S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.IC), + K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM), + IZp=ScalarDef(I32), + KZp=ScalarDef(I32), + O=TensorDef(U, + S.N, + S.OH, + S.OW, + S.IC, + S.CM, + output=True), + strides=IndexAttrDef(S.SH, + S.SW, + default=[1, 1]), + dilations=IndexAttrDef(S.DH, + S.DW, + default=[1, 1])): """Performs depth-wise 2-D convolution. Numeric casting is performed on the operands to the inner multiply, promoting @@ -482,13 +514,12 @@ @linalg_structured_op -def pooling_nhwc_sum( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.C), - K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), - O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def pooling_nhwc_sum(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): """Performs sum pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -501,13 +532,12 @@ @linalg_structured_op -def pooling_nhwc_max( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.C), - K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), - O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def pooling_nhwc_max(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): """Performs max pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -515,19 +545,21 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) - O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_signed[D.kh, D.kw]( - TypeFn.cast_signed( - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) + O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_signed[D.kh, D.kw](TypeFn.cast_signed( + U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) @linalg_structured_op -def pooling_nhwc_max_unsigned( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.C), - K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), - O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def pooling_nhwc_max_unsigned(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, + S.KH, + S.KW, + index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, S.DW, default=[1, + 1])): """Performs unsigned max pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -535,19 +567,18 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) - O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned[D.kh, D.kw]( - TypeFn.cast_unsigned( - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) + O[D.n, D.oh, D.ow, + D.c] = ReduceFn.max_unsigned[D.kh, D.kw](TypeFn.cast_unsigned( + U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) @linalg_structured_op -def pooling_nchw_max( - I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, - S.OW * S.SW + S.KW * S.DW), - K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), - O=TensorDef(U, S.N, S.C, S.OH, S.OW, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def pooling_nchw_max(I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW), + K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.C, 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 max pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -555,20 +586,17 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.c, D.oh, D.ow, D.kh, D.kw) - O[D.n, D.c, D.oh, D.ow] = ReduceFn.max_signed[D.kh, D.kw]( - TypeFn.cast_signed( - U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, - D.ow * S.SW + D.kw * S.DW,])) + O[D.n, D.c, D.oh, D.ow] = ReduceFn.max_signed[D.kh, D.kw](TypeFn.cast_signed( + U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,])) @linalg_structured_op -def pooling_nhwc_min( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.C), - K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), - O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def pooling_nhwc_min(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): """Performs min pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -576,19 +604,21 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) - O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_signed[D.kh, D.kw]( - TypeFn.cast_signed( - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) + O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_signed[D.kh, D.kw](TypeFn.cast_signed( + U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) @linalg_structured_op -def pooling_nhwc_min_unsigned( - I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, - S.C), - K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), - O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), - strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), - dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): +def pooling_nhwc_min_unsigned(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, + S.KH, + S.KW, + index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), + strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), + dilations=IndexAttrDef(S.DH, S.DW, default=[1, + 1])): """Performs unsigned min pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -596,19 +626,26 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) - O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned[D.kh, D.kw]( - TypeFn.cast_unsigned( - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) - - -@linalg_structured_op -def pooling_ndhwc_sum( - I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, - S.OW * S.SW + S.KW * S.DW, S.C), - K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]), - O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, 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])): + O[D.n, D.oh, D.ow, + D.c] = ReduceFn.min_unsigned[D.kh, D.kw](TypeFn.cast_unsigned( + U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) + + +@linalg_structured_op +def pooling_ndhwc_sum(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, + S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, + S.KD, + S.KH, + S.KW, + index_dims=[D.kd, D.kh, D.kw]), + O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, 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 3D sum pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -622,13 +659,20 @@ @linalg_structured_op -def pooling_ndhwc_max( - I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, - S.OW * S.SW + S.KW * S.DW, S.C), - K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]), - O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, 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])): +def pooling_ndhwc_max(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, + S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, + S.KD, + S.KH, + S.KW, + index_dims=[D.kd, D.kh, D.kw]), + O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, 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 3D max pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -636,20 +680,27 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw) - O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max_signed[D.kd, D.kh, D.kw]( - TypeFn.cast_signed( - U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, - D.ow * S.SW + D.kw * S.DW, D.c])) - - -@linalg_structured_op -def pooling_ndhwc_min( - I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, - S.OW * S.SW + S.KW * S.DW, S.C), - K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]), - O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, 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])): + O[D.n, D.od, D.oh, D.ow, + D.c] = ReduceFn.max_signed[D.kd, D.kh, D.kw](TypeFn.cast_signed( + U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, + D.ow * S.SW + D.kw * S.DW, D.c])) + + +@linalg_structured_op +def pooling_ndhwc_min(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, + S.OH * S.SH + S.KH * S.DH, + S.OW * S.SW + S.KW * S.DW, S.C), + K=TensorDef(T2, + S.KD, + S.KH, + S.KW, + index_dims=[D.kd, D.kh, D.kw]), + O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, 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 3D min pooling. Numeric casting is performed on the input operand, promoting it to the same @@ -657,10 +708,10 @@ """ implements(ConvolutionOpInterface) domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw) - O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min_signed[D.kd, D.kh, D.kw]( - TypeFn.cast_signed( - U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, - D.ow * S.SW + D.kw * S.DW, D.c])) + O[D.n, D.od, D.oh, D.ow, + D.c] = ReduceFn.min_signed[D.kd, D.kh, D.kw](TypeFn.cast_signed( + U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, + D.ow * S.SW + D.kw * S.DW, D.c])) @linalg_structured_op @@ -677,11 +728,10 @@ @linalg_structured_op -def fill_rng_2d( - min=ScalarDef(F64), - max=ScalarDef(F64), - seed=ScalarDef(I32), - O=TensorDef(T, S.M, S.N, output=True)): +def fill_rng_2d(min=ScalarDef(F64), + max=ScalarDef(F64), + seed=ScalarDef(I32), + O=TensorDef(T, S.M, S.N, output=True)): """Fills the output tensor with pseudo random numbers. The operation generations pseudo random numbers using a linear congruential diff --git a/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir --- a/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir +++ b/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir @@ -335,3 +335,16 @@ // CHECK: linalg.generic // CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK: = arith.subf + +// ----- + +// Verifies the fun attribute controls the binary function used. +func @generalize_copy(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { + %0 = linalg.copy ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> + return %0: tensor<4x8xf32> +} + +// CHECK-LABEL: @generalize_copy +// CHECK: linalg.generic +// CHECK-NEXT: ^bb0(%[[I:[0-9a-zA-Z]*]]: f32 +// CHECK-NEXT: linalg.yield %[[I]]