[Mlir-commits] [mlir] b22a93f - [mlir][linalg] Improve pooling op iterator order consistency
Lei Zhang
llvmlistbot at llvm.org
Tue Jan 11 09:49:42 PST 2022
Author: Lei Zhang
Date: 2022-01-11T17:49:22Z
New Revision: b22a93f4fb398c26a85b6040cdd01cc613b92c89
URL: https://github.com/llvm/llvm-project/commit/b22a93f4fb398c26a85b6040cdd01cc613b92c89
DIFF: https://github.com/llvm/llvm-project/commit/b22a93f4fb398c26a85b6040cdd01cc613b92c89.diff
LOG: [mlir][linalg] Improve pooling op iterator order consistency
All named ops list iterators for accessing output first except
pooling ops. This commit made the pooling ops consistent with
the rest.
Reviewed By: nicolasvasilache
Differential Revision: https://reviews.llvm.org/D115520
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 0dcd86a0a5e9b..dc5e5862e83c7 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -1907,18 +1907,18 @@ structured_op: !LinalgStructuredOpConfig
indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps:
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+ -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d3, d4)>
+ -> (d4, d5)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1, d2, d5)>
+ -> (d0, d1, d2, d3)>
iterator_types:
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -1978,18 +1978,18 @@ structured_op: !LinalgStructuredOpConfig
indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps:
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+ -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d3, d4)>
+ -> (d4, d5)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1, d2, d5)>
+ -> (d0, d1, d2, d3)>
iterator_types:
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -2049,18 +2049,18 @@ structured_op: !LinalgStructuredOpConfig
indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps:
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+ -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d3, d4)>
+ -> (d4, d5)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1, d2, d5)>
+ -> (d0, d1, d2, d3)>
iterator_types:
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -2191,18 +2191,18 @@ structured_op: !LinalgStructuredOpConfig
indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps:
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+ -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d3, d4)>
+ -> (d4, d5)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1, d2, d5)>
+ -> (d0, d1, d2, d3)>
iterator_types:
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -2262,18 +2262,18 @@ structured_op: !LinalgStructuredOpConfig
indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps:
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+ -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d3, d4)>
+ -> (d4, d5)>
- affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- -> (d0, d1, d2, d5)>
+ -> (d0, d1, d2, d3)>
iterator_types:
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -2336,21 +2336,21 @@ structured_op: !LinalgStructuredOpConfig
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, s12, s13] -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3
- * s10 + d6 * s12, d7)>
+ s8, s9, s10, s11, s12, s13] -> (d0, d1 * s2 + d5 * s4, d2 * s6 + d6 * s8, d3
+ * s10 + d7 * s12, d4)>
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
- s8, s9, s10, s11, s12, s13] -> (d4, d5, d6)>
+ s8, s9, s10, s11, s12, s13] -> (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, s12, s13] -> (d0, d1, d2, d3, d7)>
+ s8, s9, s10, s11, s12, s13] -> (d0, d1, d2, d3, d4)>
iterator_types:
- parallel
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -2413,21 +2413,21 @@ structured_op: !LinalgStructuredOpConfig
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, s12, s13] -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3
- * s10 + d6 * s12, d7)>
+ s8, s9, s10, s11, s12, s13] -> (d0, d1 * s2 + d5 * s4, d2 * s6 + d6 * s8, d3
+ * s10 + d7 * s12, d4)>
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
- s8, s9, s10, s11, s12, s13] -> (d4, d5, d6)>
+ s8, s9, s10, s11, s12, s13] -> (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, s12, s13] -> (d0, d1, d2, d3, d7)>
+ s8, s9, s10, s11, s12, s13] -> (d0, d1, d2, d3, d4)>
iterator_types:
- parallel
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
@@ -2490,21 +2490,21 @@ structured_op: !LinalgStructuredOpConfig
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, s12, s13] -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3
- * s10 + d6 * s12, d7)>
+ s8, s9, s10, s11, s12, s13] -> (d0, d1 * s2 + d5 * s4, d2 * s6 + d6 * s8, d3
+ * s10 + d7 * s12, d4)>
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
- s8, s9, s10, s11, s12, s13] -> (d4, d5, d6)>
+ s8, s9, s10, s11, s12, s13] -> (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, s12, s13] -> (d0, d1, d2, d3, d7)>
+ s8, s9, s10, s11, s12, s13] -> (d0, d1, d2, d3, d4)>
iterator_types:
- parallel
- parallel
- parallel
- parallel
+ - parallel
- reduction
- reduction
- reduction
- - parallel
assignments:
- !ScalarAssign
arg: O
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 9fe370ffdabb5..d3651bd766fe7 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
@@ -459,7 +459,7 @@ def pooling_nhwc_sum(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+ domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.c] += TypeFn.cast(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
@@ -478,7 +478,7 @@ def pooling_nhwc_max(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+ domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.c] = ReduceFn.max[D.kh, D.kw](
TypeFn.cast(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
@@ -498,7 +498,7 @@ def pooling_nhwc_max_unsigned(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+ 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]))
@@ -539,7 +539,7 @@ def pooling_nhwc_min(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+ domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.c] = ReduceFn.min[D.kh, D.kw](
TypeFn.cast(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
@@ -559,7 +559,7 @@ def pooling_nhwc_min_unsigned(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+ 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]))
@@ -579,7 +579,7 @@ def pooling_ndhwc_sum(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
+ 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] += TypeFn.cast(
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])
@@ -599,7 +599,7 @@ def pooling_ndhwc_max(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
+ 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[D.kd, D.kh, D.kw](
TypeFn.cast(
U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
@@ -620,7 +620,7 @@ def pooling_ndhwc_min(
data type as the accumulator/output.
"""
implements(ConvolutionOpInterface)
- domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
+ 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[D.kd, D.kh, D.kw](
TypeFn.cast(
U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
More information about the Mlir-commits
mailing list