[mlir][linalg][python] Add attribute support to the YAML codegen.

Extend the yaml code generation to support the index attributes that https://reviews.llvm.org/D104711 added to the OpDSL. Differential Revision: https://reviews.llvm.org/D104712
llvm · Jun 24, 2021 · 25bb616 · 25bb616
1 parent adace79
commit 25bb616
Show file tree

Hide file tree

Showing 6 changed files with 390 additions and 100 deletions.
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -13,19 +13,19 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: input
-    shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
+    usage: InputOperand
     type_var: T1
+    shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
   - !LinalgOperandDefConfig
     name: B
-    usage: input
-    shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
+    usage: InputOperand
     type_var: T2
+    shape_map: affine_map<()[s0, s1, s2] -> (s2, s1)>
   - !LinalgOperandDefConfig
     name: C
-    usage: output
-    shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
+    usage: OutputOperand
     type_var: U
+    shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
@@ -75,19 +75,19 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: input
-    shape: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s3)>
+    usage: InputOperand
     type_var: T1
+    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s3)>
   - !LinalgOperandDefConfig
     name: B
-    usage: input
-    shape: affine_map<()[s0, s1, s2, s3] -> (s0, s3, s2)>
+    usage: InputOperand
     type_var: T2
+    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s3, s2)>
   - !LinalgOperandDefConfig
     name: C
-    usage: output
-    shape: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s2)>
+    usage: OutputOperand
     type_var: U
+    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s2)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1, d3)>
@@ -138,19 +138,19 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: input
-    shape: affine_map<()[s0, s1] -> (s0, s1)>
+    usage: InputOperand
     type_var: T1
+    shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: y
-    usage: input
-    shape: affine_map<()[s0, s1] -> (s1)>
+    usage: InputOperand
     type_var: T2
+    shape_map: affine_map<()[s0, s1] -> (s1)>
   - !LinalgOperandDefConfig
     name: x
-    usage: output
-    shape: affine_map<()[s0, s1] -> (s0)>
+    usage: OutputOperand
     type_var: U
+    shape_map: affine_map<()[s0, s1] -> (s0)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
@@ -199,19 +199,19 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: y
-    usage: input
-    shape: affine_map<()[s0, s1] -> (s1)>
+    usage: InputOperand
     type_var: T1
+    shape_map: affine_map<()[s0, s1] -> (s1)>
   - !LinalgOperandDefConfig
     name: A
-    usage: input
-    shape: affine_map<()[s0, s1] -> (s1, s0)>
+    usage: InputOperand
     type_var: T2
+    shape_map: affine_map<()[s0, s1] -> (s1, s0)>
   - !LinalgOperandDefConfig
     name: x
-    usage: output
-    shape: affine_map<()[s0, s1] -> (s0)>
+    usage: OutputOperand
     type_var: U
+    shape_map: affine_map<()[s0, s1] -> (s0)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d1)>
@@ -260,19 +260,19 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: input
-    shape: affine_map<()[s0] -> (s0)>
+    usage: InputOperand
     type_var: T1
+    shape_map: affine_map<()[s0] -> (s0)>
   - !LinalgOperandDefConfig
     name: B
-    usage: input
-    shape: affine_map<()[s0] -> (s0)>
+    usage: InputOperand
     type_var: T2
+    shape_map: affine_map<()[s0] -> (s0)>
   - !LinalgOperandDefConfig
     name: C
-    usage: output
-    shape: affine_map<()[s0] -> ()>
+    usage: OutputOperand
     type_var: U
+    shape_map: affine_map<()[s0] -> ()>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0)[s0] -> (d0)>
@@ -306,6 +306,83 @@ structured_op: !LinalgStructuredOpConfig
                 - !ScalarExpression
                   scalar_arg: B
 --- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+  name: depthwise_conv_2d_input_nhwc_filter_hwc_poly
+  cpp_class_name: DepthwiseConv2DInputNhwcFilterHwcPolyOp
+  doc: A depth-wise 2-D convolution operation.
+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] ->
+      (s0, s6, s7, 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] ->
+      (s4, s5, s3)>
+  - !LinalgOperandDefConfig
+    name: O
+    usage: OutputOperand
+    type_var: U
+    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
+      (s0, s1, s2, s3)>
+  - !LinalgOperandDefConfig
+    name: strides
+    usage: IndexAttribute
+    type_var: I64
+    attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
+      -> (s8, s9)>
+  - !LinalgOperandDefConfig
+    name: dilations
+    usage: IndexAttribute
+    type_var: I64
+    attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
+      -> (s10, s11)>
+  indexing_maps: !LinalgIndexingMapsConfig
+    static_indexing_maps:
+    - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9,
+      s10, s11] -> (d0, d1 * s8 + d4 * s10, d2 * s9 + d5 * s11, d3)>
+    - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9,
+      s10, s11] -> (d4, d5, d3)>
+    - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9,
+      s10, s11] -> (d0, d1, d2, d3)>
+  iterator_types:
+  - parallel
+  - parallel
+  - parallel
+  - parallel
+  - 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: fill_rng_2d
   cpp_class_name: FillRng2DOp
@@ -323,21 +400,21 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: min
-    usage: input
+    usage: InputOperand
     type_var: F64
   - !LinalgOperandDefConfig
     name: max
-    usage: input
+    usage: InputOperand
     type_var: F64
   - !LinalgOperandDefConfig
     name: seed
-    usage: input
+    usage: InputOperand
     type_var: I32
   - !LinalgOperandDefConfig
     name: O
-    usage: output
-    shape: affine_map<()[s0, s1] -> (s0, s1)>
+    usage: OutputOperand
     type_var: T
+    shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> ()>

diff --git a/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir
@@ -30,6 +30,36 @@ func @generalize_matmul_tensor_i32(%A : tensor<16x8xi32>, %B: tensor<8x32xi32>,
 
 // -----
 
+func @generalize_depthwise_conv_2d_input_nhwc_filter_hwc_poly_f32(%input : tensor<1x4x16x1xf32>, %filter: tensor<2x2x1xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> {
+  %0 = linalg.depthwise_conv_2d_input_nhwc_filter_hwc_poly {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
+    ins(%input, %filter : tensor<1x4x16x1xf32>, tensor<2x2x1xf32>) outs(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
+  return %0: tensor<1x2x4x1xf32>
+}
+
+// CHECK-LABEL: @generalize_depthwise_conv_2d_input_nhwc_filter_hwc_poly_f32
+// CHECK:      ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[FILTER_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32)
+// CHECK-NEXT:   %[[MUL:.+]] = mulf %[[IN_ARG]], %[[FILTER_ARG]] : f32
+// CHECK-NEXT:   %[[ADD:.+]] = addf %[[OUT_ARG]], %[[MUL]] : f32
+// CHECK-NEXT:   linalg.yield %[[ADD]] : f32
+// CHECK-NEXT: -> tensor<1x2x4x1xf32>
+
+// -----
+
+func @generalize_depthwise_conv_2d_input_nhwc_filter_hwc_poly_i32(%input : tensor<1x4x16x1xi32>, %filter: tensor<2x2x1xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> {
+  %0 = linalg.depthwise_conv_2d_input_nhwc_filter_hwc_poly {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
+    ins(%input, %filter : tensor<1x4x16x1xi32>, tensor<2x2x1xi32>) outs(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32>
+  return %0: tensor<1x2x4x1xi32>
+}
+
+// CHECK-LABEL: @generalize_depthwise_conv_2d_input_nhwc_filter_hwc_poly_i32
+// CHECK:      ^{{.*}}(%[[IN_ARG:.+]]: i32, %[[FILTER_ARG:.+]]: i32, %[[OUT_ARG:.+]]: i32)
+// CHECK-NEXT:   %[[MUL:.+]] = muli %[[IN_ARG]], %[[FILTER_ARG]] : i32
+// CHECK-NEXT:   %[[ADD:.+]] = addi %[[OUT_ARG]], %[[MUL]] : i32
+// CHECK-NEXT:   linalg.yield %[[ADD]] : i32
+// CHECK-NEXT: -> tensor<1x2x4x1xi32>
+
+// -----
+
 func @generalize_fill_rng_2d_f32(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> {
   %0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) outs(%O : tensor<16x32xf32>) -> tensor<16x32xf32>
   return %0: tensor<16x32xf32>

diff --git a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml
@@ -21,9 +21,9 @@ structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: O
-    usage: output
-    shape: affine_map<()[s0, s1] -> (s0, s1)>
+    usage: OutputOperand
     type_var: T
+    shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
@@ -86,12 +86,13 @@ structured_op: !LinalgStructuredOpConfig
 
 # @linalg_structured_op
 # def test2(I=TensorDef(T, S.M, S.N),
-#           O=TensorDef(T, S.M, S.N, output=True)):
+#           O=TensorDef(T, S.M, S.N, output=True),
+#           strides=AttributeDef(S.SM, S.SN)):
 #   """Title.
 
 #   Detailed description.
 #   """
-#   O[D.m, D.n] = I[D.n, D.m]
+#   O[D.m, D.n] = I[D.n * S.SM, D.m * S.SN]
 
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
@@ -103,49 +104,67 @@ metadata: !LinalgOpMetadata
     Detailed description.
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !LinalgOperandDefConfig
-    name: value
-    usage: input
-    type_var: T
   - !LinalgOperandDefConfig
     name: I
-    usage: input
-    shape: affine_map<()[s0, s1] -> (s1, s0)>
+    usage: InputOperand
     type_var: T
+    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: O
-    usage: output
-    shape: affine_map<()[s0, s1] -> (s0, s1)>
+    usage: OutputOperand
     type_var: T
+    shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
+  - !LinalgOperandDefConfig
+    name: strides
+    usage: IndexAttribute
+    type_var: I64
+    attribute_map: affine_map<()[s0, s1, s2, s3] -> (s2, s3)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
-    - affine_map<(d0, d1)[s0, s1] -> ()>
-    - affine_map<(d0, d1)[s0, s1] -> (d1, d0)>
-    - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
+    - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>
+    - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      scalar_apply:
-        fn_name: add
-        operands:
-        - !ScalarExpression
-          scalar_arg: value
-        - !ScalarExpression
-          scalar_arg: I
+      scalar_arg: I
 
-# IMPL-LABEL:  Test2Op::iterator_types()
-#  IMPL-NEXT:  { getParallelIteratorTypeName(), getParallelIteratorTypeName() }
+# ODS-LABEL:  def Test2Op : LinalgStructuredBase_Op<"test2"
+
+#       ODS:  let arguments =
+#  ODS-NEXT:    Variadic<AnyType>:$inputs,
+#  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
+#  ODS-NEXT:    RankedI64ElementsAttr<[2]>:$strides
+
+#       ODS:  "Attribute":$strides
+#       ODS:  $_state.addAttribute("strides", strides);
+
+#       ODS:  bool hasDynamicIndexingMaps();
+#  ODS-NEXT:  LogicalResult verifyIndexingMapRequiredAttributes();
+
+#       IMPL:  getSymbolBindings(Test2Op self)
+#       IMPL:  cst2 = self.strides().getValue<int64_t>({ 0 });
+#  IMPL-NEXT:  getAffineConstantExpr(cst2, context)
+#       IMPL:  cst3 = self.strides().getValue<int64_t>({ 1 });
+#  IMPL-NEXT:  getAffineConstantExpr(cst3, context)
 
 #       IMPL:  Test2Op::indexing_maps()
-#       IMPL:  "affine_map<(d0, d1)[s0, s1] -> ()>"
-#       IMPL:  "affine_map<(d0, d1)[s0, s1] -> (d1, d0)>"
-#       IMPL:  "affine_map<(d0, d1)[s0, s1] -> (d0, d1)>"
+#       IMPL:  = getSymbolBindings(*this);
+#       IMPL:  "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>"
+#       IMPL:  "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>"
+
+#       IMPL:  Test2Op::getNumRegionArgs() { return 2; }
+
+#       IMPL:  Test2Op::hasDynamicIndexingMaps() { return true; }
+#       IMPL:  Test2Op::verifyIndexingMapRequiredAttributes()
+#       IMPL:  auto attr = op->getAttrOfType<DenseElementsAttr>("strides")
+#       IMPL:  "missing indexing map required attribute 'strides'"
 
 #       IMPL:  void Test2Op::regionBuilder(
-#       IMPL:    ImplicitLocOpBuilder &b, Block &block, ValueRange captures)
+#  IMPL-NEXT:    ImplicitLocOpBuilder &b, Block &block, ValueRange captures)
+#  IMPL-NEXT:    assert(2 > 0 && block.getNumArguments() == 2 &&
 
-#       IMPL:   = helper.applyfn__add(block.getArgument(0), block.getArgument(1));
+#       IMPL:   yields.push_back(block.getArgument(0));