Further fixes per @banach-space's comments

Author: rolfmorel

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

@@ -698,9 +698,9 @@ def ContractOp : LinalgStructuredBase_Op<"contract", [
 
     where `I`, `J`, and `H` are multi-indices, i.e. sequences/ordered sets of
     dimension identifiers (meant to range over valid indices), corresponding to
-    the co-domains of the (projected permutation) `indexing_maps` of `A`, `B`
-    and `C`, respectively. `SUM_{dims}` means reduce over all valid indices for
-    the dimensions in the set `dims`.
+    the co-domains of the mandatory (projected permutation) `indexing_maps` of
+    `A`, `B` and `C`, respectively. `SUM_{dims}` means reduce over all valid
+    indices for the dimensions in the set `dims`.
 
     The iteration space consists of all dimensions in `I`, `J` and `H`, i.e. the
     domain of each of the `affine_map`s. Like for einsums, the iteration type of
@@ -719,21 +719,24 @@ def ContractOp : LinalgStructuredBase_Op<"contract", [
     `n` and `b` are of parallel iteration-type) and gets represented as:
 
     ```
-    %0 = linalg.contract
+    %D = linalg.contract
         indexing_maps = [affine_map<(batch, m, n, k) -> (batch, m, k)>,
                          affine_map<(batch, m, n, k) -> (batch, k, n)>,
                          affine_map<(batch, m, n, k) -> (batch, m, n)>]
-        ins(%arg0, %arg1: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
-        outs(%arg2: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+        ins(%A, %B: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+        outs(%C: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
     ```
 
     Note that by permuting the dims in the co-domains of the `affine_map`s, we
     can apply arbitrary transposes to the inputs and output. Similarly,
     arbitrary broadcasts can be achieved through leaving out dims on either
-    input operand.
+    input operand - these dims' inferred iter type will be parallel.
 
     Numeric casting is performed on the operands to the inner multiplication,
     promoting them to the same data type as the accumulator/output.
+
+    TODO: Allow control over the combining/accumulating op and possibly the
+          multiplication op.
   }];
 
   let arguments = (ins

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -3693,10 +3693,18 @@ void ContractOp::print(OpAsmPrinter &p) {
 
 LogicalResult ContractOp::verify() {
   int iterationSpaceDims = -1;
-  // Maps iter space dim (as index) to num of occurrences in inputs and output.
+  // Map iter space dims to #occurrences in inputs' and output's affine_maps:
+  // e.g., inOccurrences[0] will hold #times that dim (with index) 0 is used to
+  // access an input operand (so occurrence count can be at most 2) and
+  // outOccurrences[1] will indicate whether dim 1 occurred in the output, etc.
   SmallVector<size_t> inOccurrences;
   SmallVector<size_t> outOccurrences;
 
+  // For each operand's affine_map and type, check that the rank of the
+  // affine_map's domain is the same as those seen prior, check that the
+  // affine_map's co-domain rank is the same as that of the corresponding type,
+  // check that the affine_map is a projected permutation, and, finally, update
+  // inputs and output occurrence counts for dims in the co-domains.
   auto checkAffineMapAndType = [&](AffineMap affineMap, Type operandType,
                                    bool isInput) -> LogicalResult {
     if (iterationSpaceDims == -1) {
@@ -3737,14 +3745,15 @@ LogicalResult ContractOp::verify() {
        llvm::zip(getIndexingMapsArray(), getOperandTypes(),
                  SmallVector<bool>{true, true, false})) {
     if (failed(checkAffineMapAndType(affineMap, operandType, isInput)))
-      return failure(); // NOTE: checking lambda will emit error.
+      return failure(); // NB: checkAffineMapAndType will emit relevant error.
   }
 
   bool hasContractingDim = false;
   for (size_t dimIndex = 0; dimIndex < (size_t)iterationSpaceDims; dimIndex++) {
     size_t inOccCount = inOccurrences[dimIndex];
     size_t outOccCount = outOccurrences[dimIndex];
 
+    // We have a contracting dim if and only if ...
     hasContractingDim |= inOccCount == 2 && outOccCount == 0;
 
     if (inOccCount == 0 && outOccCount == 0)

mlir/test/Dialect/Linalg/generalize-named-ops.mlir

@@ -1016,13 +1016,13 @@ func.func @matmul_transpose_a_b_explicit(%arg0: memref<5x3xf32>, %arg1: memref<7
 // CHECK-NEXT:      arith.addf
 // CHECK-NEXT:      linalg.yield
 
-func.func @contract_matmul(%arg0: memref<3x5xf32>, %arg1: memref<5x7xf32>, %arg2: memref<3x7xf32>) {
+func.func @contract_matmul(%A: memref<3x5xf32>, %B: memref<5x7xf32>, %C: memref<3x7xf32>) {
   linalg.contract
       indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
                        affine_map<(d0, d1, d2) -> (d2, d1)>,
                        affine_map<(d0, d1, d2) -> (d0, d1)>]
-      ins(%arg0, %arg1 : memref<3x5xf32>, memref<5x7xf32>)
-      outs(%arg2: memref<3x7xf32>)
+      ins(%A, %B : memref<3x5xf32>, memref<5x7xf32>)
+      outs(%C: memref<3x7xf32>)
 
   return
 }
@@ -1046,13 +1046,13 @@ func.func @contract_matmul(%arg0: memref<3x5xf32>, %arg1: memref<5x7xf32>, %arg2
 // CHECK-NEXT:      arith.addf
 // CHECK-NEXT:      linalg.yield
 
-func.func @contract_matmul_transpose_a_b(%arg0: memref<5x3xf32>, %arg1: memref<7x5xf32>, %arg2: memref<3x7xf32>) {
+func.func @contract_matmul_transpose_a_b(%A: memref<5x3xf32>, %B: memref<7x5xf32>, %C: memref<3x7xf32>) {
   linalg.contract
       indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0)>,
                        affine_map<(d0, d1, d2) -> (d1, d2)>,
                        affine_map<(d0, d1, d2) -> (d0, d1)>]
-      ins(%arg0, %arg1 : memref<5x3xf32>, memref<7x5xf32>)
-      outs(%arg2: memref<3x7xf32>)
+      ins(%A, %B : memref<5x3xf32>, memref<7x5xf32>)
+      outs(%C: memref<3x7xf32>)
   return
 }
 
@@ -1075,13 +1075,13 @@ func.func @contract_matmul_transpose_a_b(%arg0: memref<5x3xf32>, %arg1: memref<7
 // CHECK-NEXT:      arith.addf
 // CHECK-NEXT:      linalg.yield
 
-func.func @contract_batch_matmul(%arg0: memref<9x3x5xf32>, %arg1: memref<9x5x7xf32>, %arg2: memref<9x3x7xf32>) {
+func.func @contract_batch_matmul(%A: memref<9x3x5xf32>, %B: memref<9x5x7xf32>, %C: memref<9x3x7xf32>) {
   linalg.contract
       indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
                        affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]
-      ins(%arg0, %arg1 : memref<9x3x5xf32>, memref<9x5x7xf32>)
-      outs(%arg2: memref<9x3x7xf32>)
+      ins(%A, %B : memref<9x3x5xf32>, memref<9x5x7xf32>)
+      outs(%C: memref<9x3x7xf32>)
   return
 }
 

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

@@ -121,7 +121,7 @@ func.func @generalize_matmul_unsigned_tensor_f16f64i32(%A : tensor<16x8xf16>, %B
 // -----
 
 func.func @generalize_matmul_as_contraction_tensor_f16f64f32(
-    %A : tensor<16x8xf16>,
+    %A: tensor<16x8xf16>,
     %B: tensor<8x32xf64>,
     %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
   %0 = linalg.contract

mlir/test/Dialect/Linalg/invalid.mlir

@@ -549,7 +549,8 @@ func.func @invalid_indexing_maps_placement_matmul(%lhs: tensor<4x1xf32>, %rhs: t
 
 func.func @invalid_indexing_maps_placement_contraction(
     %lhs: tensor<4x1xf32>, %rhs: tensor<1x64xf32>, %init: tensor<4x64xf32>) {
-  // expected-error @+2 {{custom op 'linalg.contract' expected 'indexing_maps' attribute}}
+  // expected-error @+3 {{custom op 'linalg.contract' expected 'indexing_maps' attribute}}
+  // NB: indexing_maps should be provided before ins and outs
   linalg.contract
       ins(%lhs, %rhs : tensor<4x1xf32>, tensor<1x64xf32>)
       outs(%init : tensor<4x64xf32>)

mlir/test/Dialect/Linalg/named-ops.mlir

@@ -1568,7 +1568,6 @@ func.func @contract_matmul_bcast_b(%A: memref<3x5xf32>, %B: memref<5xf32>, %C: m
 
 // -----
 
-
 // CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d2)>
 // CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 // CHECK-LABEL: func.func @contract_matmul_bcast_a_b
@@ -1608,7 +1607,6 @@ func.func @contract_matmul_bcast_a_transpose_b(
   return
 }
 
-
 // -----
 
 // CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
@@ -1629,7 +1627,6 @@ func.func @contract_matmul_bcast_b_transpose_a(%A: memref<5x3xf32>, %B: memref<5
   return
 }
 
-
 // -----
 
 // CHECK-LABEL: func @mmt4d