[mlir][linalg] Add one-shot-bufferize tests for Linalg ops: reduce, map and transpose.

[pifon2a]

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

350

so, if the pass did not bufferize these ops at all, the CHECKs would still match it. Maybe add a memref type somewhere, like in line 326?

[olegshyshkov]

Update tests to check for memrefs.

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

let extraClassDeclaration = structuredOpsBaseDecls # [{
     std::pair<int64_t, int64_t> getOutputsPositionRange() {
       int64_t getNumOperands = this->getNumOperands();
       return {getNumOperands - 1, getNumOperands};
     }
     OpOperandVector getOpOperandsMatchingBBargs() {
       return getInputOperands();
     }
 
+    bool payloadUsesValueFromOperand(OpOperand * opOperand) {
+      if (isOutput(opOperand)) return false;
+      return !getMatchingBlockArgument(opOperand).use_empty();
+    }
+
     static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
                               mlir::ArrayRef<mlir::NamedAttribute>)>
     getRegionBuilder() {
       return nullptr;
     }
   }];
 
   let hasCustomAssemblyFormat = 1;

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

%r1 = linalg.generic #trait outs (%r0 : tensor<?xf32>) {
     } -> (tensor<?xf32>)
 
   // CHECK: return %[[ALLOC]]
   return %r1 : tensor<?xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @map_binary
+// CHECK-SAME:  %[[LHS:[0-9a-zA-Z]*]]: memref<64xf32
+// CHECK-SAME:  %[[RHS:[0-9a-zA-Z]*]]: memref<64xf32
+func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
+                      %init: tensor<64xf32>) -> tensor<64xf32> {
+   // CHECK:      linalg.map
+   // CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<64xf32
+   %add = linalg.map
+          ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>)
+          outs(%init:tensor<64xf32>)
+          (%lhs_elem: f32, %rhs_elem: f32) {
+            %0 = arith.addf %lhs_elem, %rhs_elem: f32
+            linalg.yield %0: f32
+          }
+  func.return %add : tensor<64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @reduce
+// CHECK-SAME:  %[[INPUT:.*]]: memref<16x32x64xf32
+func.func @reduce(%input: tensor<16x32x64xf32>,
+                  %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
+  // CHECK:     linalg.reduce
+  // CHECK-SAME: ins(%[[INPUT]] : memref<16x32x64xf32
+  %reduce = linalg.reduce
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<16x64xf32>)
+      dimensions = [1]
+      (%in: f32, %out: f32) {
+        %0 = arith.addf %in, %out: f32
+        linalg.yield %0: f32
+      }
+  func.return %reduce : tensor<16x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @transpose
+// CHECK-SAME:  %[[ARG0:.*]]: memref<16x32x64xf32
+func.func @transpose(%input: tensor<16x32x64xf32>,
+                     %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
+  // CHECK:      linalg.transpose
+  // CHECK-SAME: ins(%[[ARG0]] : memref<16x32x64xf32
+  %transpose = linalg.transpose
+      ins(%input:tensor<16x32x64xf32>)
+      outs(%init:tensor<32x64x16xf32>)
+      permutation = [1, 2, 0]
+  func.return %transpose : tensor<32x64x16xf32>
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // AllocTensorOp elimination would produce SSA violations for the example below.
 //===----------------------------------------------------------------------===//
 
 func.func @depthwise_conv_1d_nwc_wc(%arg0: index, %arg1: index, %arg2: tensor<8x18x32xf32>)
     -> tensor<?x1x6x8xf32> {
   %c0 = arith.constant 0 : index
   %c32 = arith.constant 32 : index