[Mlir-commits] [mlir] a8de667 - [mlir] Add op for NCHW conv2d.
Stella Laurenzo
llvmlistbot at llvm.org
Sun Aug 22 17:29:16 PDT 2021
Author: Stella Laurenzo
Date: 2021-08-22T17:27:33-07:00
New Revision: a8de667af092c9b4b3b4a95827a521602ebf14ed
URL: https://github.com/llvm/llvm-project/commit/a8de667af092c9b4b3b4a95827a521602ebf14ed
DIFF: https://github.com/llvm/llvm-project/commit/a8de667af092c9b4b3b4a95827a521602ebf14ed.diff
LOG: [mlir] Add op for NCHW conv2d.
* This is the native data layout for PyTorch and npcomp was using the prior version before cleanup.
Differential Revision: https://reviews.llvm.org/D108527
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 5ce6d4c57e460..ec71aa59ab916 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -895,6 +895,10 @@ metadata: !LinalgOpMetadata
doc: |-
Performs 2-D convolution.
+ Layout:
+ * Input: NHWC.
+ * Kernel: HWCF.
+
Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.
structured_op: !LinalgStructuredOpConfig
@@ -977,6 +981,10 @@ metadata: !LinalgOpMetadata
doc: |-
Performs 2-D convolution with zero point offsets.
+ Layout:
+ * Input: NHWC.
+ * Kernel: HWCF.
+
Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. This includes the zero
point offsets common to quantized operations.
@@ -1086,6 +1094,92 @@ structured_op: !LinalgStructuredOpConfig
- !ScalarExpression
scalar_arg: KZp
--- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+ name: conv_2d_nchw_fchw
+ cpp_class_name: Conv2DNchwFchwOp
+ doc: |-
+ Performs 2-D convolution.
+
+ Layout:
+ * Input: NCHW.
+ * Kernel: FCHW.
+
+ Numeric casting is performed on the operands to the inner multiply, promoting
+ them to the same data type as the accumulator/output.
+structured_op: !LinalgStructuredOpConfig
+ args:
+ - !LinalgOperandDefConfig
+ name: I
+ usage: InputOperand
+ type_var: T1
+ shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12]
+ -> (s0, s1, s2, s3)>
+ - !LinalgOperandDefConfig
+ name: K
+ usage: InputOperand
+ type_var: T2
+ shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12]
+ -> (s4, s1, s5, s6)>
+ - !LinalgOperandDefConfig
+ name: O
+ usage: OutputOperand
+ type_var: U
+ shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12]
+ -> (s0, s4, s7, s8)>
+ - !LinalgOperandDefConfig
+ name: strides
+ usage: IndexAttribute
+ type_var: I64
+ attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
+ s12] -> (s9, s10)>
+ - !LinalgOperandDefConfig
+ name: dilations
+ usage: IndexAttribute
+ type_var: I64
+ attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
+ s12] -> (s11, s12)>
+ indexing_maps: !LinalgIndexingMapsConfig
+ static_indexing_maps:
+ - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
+ s9, s10, s11, s12] -> (d0, d4, d2 * s9 + d5 * s11, d3 * s10 + d6 * s12)>
+ - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
+ s9, s10, s11, s12] -> (d1, d4, d5, d6)>
+ - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
+ s9, s10, s11, s12] -> (d0, d1, d2, d3)>
+ iterator_types:
+ - parallel
+ - parallel
+ - parallel
+ - parallel
+ - reduction
+ - reduction
+ - reduction
+ assignments:
+ - !ScalarAssign
+ arg: O
+ value: !ScalarExpression
+ scalar_apply:
+ fn_name: add
+ operands:
+ - !ScalarExpression
+ scalar_arg: O
+ - !ScalarExpression
+ scalar_apply:
+ fn_name: mul
+ operands:
+ - !ScalarExpression
+ symbolic_cast:
+ type_var: U
+ operands:
+ - !ScalarExpression
+ scalar_arg: I
+ - !ScalarExpression
+ symbolic_cast:
+ type_var: U
+ operands:
+ - !ScalarExpression
+ scalar_arg: K
+--- !LinalgOpConfig
metadata: !LinalgOpMetadata
name: conv_3d_ndhwc_dhwcf
cpp_class_name: Conv3DNdhwcDhwcfOp
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
index 29b1397d1c20a..38db294428c94 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
@@ -212,6 +212,10 @@ def conv_2d_nhwc_hwcf(
dilations=AttributeDef(S.DH, S.DW)):
"""Performs 2-D convolution.
+ Layout:
+ * Input: NHWC.
+ * Kernel: HWCF.
+
Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.
"""
@@ -231,6 +235,10 @@ def conv_2d_nhwc_hwcf_q(
dilations=AttributeDef(S.DH, S.DW)):
"""Performs 2-D convolution with zero point offsets.
+ Layout:
+ * Input: NHWC.
+ * Kernel: HWCF.
+
Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. This includes the zero
point offsets common to quantized operations.
@@ -240,6 +248,27 @@ def conv_2d_nhwc_hwcf_q(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c
]) - cast(U, IZp)) * (cast(U, K[D.kh, D.kw, D.c, D.f]) - cast(U, KZp))
+ at linalg_structured_op
+def conv_2d_nchw_fchw(
+ I=TensorDef(T1, S.N, S.C, S.IH, S.IW),
+ 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=AttributeDef(S.SH, S.SW),
+ dilations=AttributeDef(S.DH, S.DW)):
+ """Performs 2-D convolution.
+
+ Layout:
+ * Input: NCHW.
+ * Kernel: FCHW.
+
+ Numeric casting is performed on the operands to the inner multiply, promoting
+ them to the same data type as the accumulator/output.
+ """
+ domain(D.n, D.f, D.oh, D.ow, D.c, D.kh, D.kw)
+ O[D.n, D.f, D.oh, D.ow] += cast(
+ U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW
+ ]) * cast(U, K[D.f, D.c, D.kh, D.kw])
+
@linalg_structured_op
def conv_3d_ndhwc_dhwcf(
I=TensorDef(T1, S.N, S.ID, S.IH, S.IW, S.C),
More information about the Mlir-commits
mailing list