[FoldSCFIf] added more test cases

Author: kumasento

tools/polymer/lib/Transforms/FoldSCFIf.cc

@@ -167,9 +167,8 @@ struct MemRefStoreInfo {
 };
 } // namespace
 
-static void
-getMemRefStoreInfo(Block *block,
-                   SmallDenseMap<Value, MemRefStoreInfo> &storeInfo) {
+static void getMemRefStoreInfo(Block *block,
+                               MapVector<Value, MemRefStoreInfo> &storeInfo) {
   unsigned ord = 0;
   for (Operation &op : block->getOperations())
     if (isa<memref::StoreOp, mlir::AffineStoreOp>(op)) {
@@ -193,7 +192,7 @@ static LogicalResult liftStoreOps(scf::IfOp ifOp, FuncOp f, OpBuilder &b) {
   if (!hasMatchingStores({ifOp.thenBlock(), ifOp.elseBlock()}))
     return failure();
 
-  SmallDenseMap<Value, MemRefStoreInfo> storeInfo;
+  MapVector<Value, MemRefStoreInfo> storeInfo;
   getMemRefStoreInfo(ifOp.thenBlock(), storeInfo);
 
   SmallVector<Type> storeTypes(storeInfo.size());
@@ -244,6 +243,7 @@ static LogicalResult liftStoreOps(scf::IfOp ifOp, FuncOp f, OpBuilder &b) {
   cloneBlock(newIfOp.elseBlock(), ifOp.elseBlock());
 
   b.setInsertionPointAfter(newIfOp);
+
   for (auto &p : storeInfo) {
     Value memref;
     MemRefStoreInfo info;

tools/polymer/test/polymer-opt/FoldSCFIf/match-store-addr.mlir

@@ -0,0 +1,15 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%A: memref<10xf32>, %i: index, %a: f32, %b: f32, %cond: i1) {
+  scf.if %cond {
+    affine.store %a, %A[%i] : memref<10xf32>
+  } else {
+    affine.store %b, %A[%i] : memref<10xf32>
+  }
+
+  return
+}
+
+// CHECK: func @foo(%[[A:.*]]: memref<10xf32>, %[[i:.*]]: index, %[[a:.*]]: f32, %[[b:.*]]: f32, %[[cond:.*]]: i1) 
+// CHECK-NEXT:   %[[v0:.*]] = select %[[cond]], %[[a]], %[[b]] : f32
+// CHECK-NEXT:   affine.store %[[v0]], %[[A]][%[[i]]] : memref<10xf32>

tools/polymer/test/polymer-opt/FoldSCFIf/match-store-diff-mem.mlir

@@ -0,0 +1,18 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%A: memref<10xf32>, %B: memref<10xf32>, %a: f32, %b: f32, %cond: i1) {
+  scf.if %cond {
+    affine.store %a, %A[0] : memref<10xf32>
+  } else {
+    affine.store %b, %B[0] : memref<10xf32>
+  }
+  return
+}
+
+// CHECK: func @foo(%[[A:.*]]: memref<10xf32>, %[[B:.*]]: memref<10xf32>, %[[a:.*]]: f32, %[[b:.*]]: f32, %[[cond:.*]]: i1) 
+// CHECK-NEXT:   %[[v0:.*]] = affine.load %[[B]][0] : memref<10xf32>
+// CHECK-NEXT:   %[[v1:.*]] = affine.load %[[A]][0] : memref<10xf32>
+// CHECK-NEXT:   %[[v2:.*]] = select %[[cond]], %[[a]], %[[v1]] : f32
+// CHECK-NEXT:   %[[v3:.*]] = select %[[cond]], %[[v0]], %[[b]] : f32
+// CHECK-NEXT:   affine.store %[[v2]], %[[A]][0] : memref<10xf32>
+// CHECK-NEXT:   affine.store %[[v3]], %[[B]][0] : memref<10xf32>

tools/polymer/test/polymer-opt/FoldSCFIf/match-store-no-else.mlir

@@ -0,0 +1,13 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%A: memref<10xf32>, %a: f32, %cond: i1) {
+  scf.if %cond {
+    affine.store %a, %A[0] : memref<10xf32>
+  }
+  return
+}
+
+// CHECK: func @foo(%[[A:.*]]: memref<10xf32>, %[[a:.*]]: f32, %[[cond:.*]]: i1) 
+// CHECK-NEXT:   %[[v0:.*]] = affine.load %[[A]][0] : memref<10xf32>
+// CHECK-NEXT:   %[[v1:.*]] = select %[[cond]], %[[a]], %[[v0]] : f32
+// CHECK-NEXT:   affine.store %[[v1]], %[[A]][0] : memref<10xf32>

tools/polymer/test/polymer-opt/FoldSCFIf/match-store.mlir

@@ -1,13 +1,14 @@
 // RUN: polymer-opt %s -fold-scf-if | FileCheck %s
 
-func @foo(%A: memref<10xf32>, %a: f32, %cond: i1) {
+func @foo(%A: memref<10xf32>, %a: f32, %b: f32, %cond: i1) {
   scf.if %cond {
     affine.store %a, %A[0] : memref<10xf32>
+  } else {
+    affine.store %b, %A[0] : memref<10xf32>
   }
   return
 }
 
-// CHECK: func @foo(%[[A:.*]]: memref<10xf32>, %[[a:.*]]: f32, %[[cond:.*]]: i1) 
-// CHECK-NEXT:   %[[v0:.*]] = affine.load %[[A]][0] : memref<10xf32>
-// CHECK-NEXT:   %[[v1:.*]] = select %[[cond]], %[[a]], %[[v0]] : f32
-// CHECK-NEXT:   affine.store %[[v1]], %[[A]][0] : memref<10xf32>
+// CHECK: func @foo(%[[A:.*]]: memref<10xf32>, %[[a:.*]]: f32, %[[b:.*]]: f32, %[[cond:.*]]: i1) 
+// CHECK-NEXT:   %[[v0:.*]] = select %[[cond]], %[[a]], %[[b]] : f32
+// CHECK-NEXT:   affine.store %[[v0]], %[[A]][0] : memref<10xf32>