[mlir][SCF][bufferize] Bufferize scf.if/execute_region terminators separately

[springerm]

Summary

This allows for better type inference during bufferization and is in preparation of supporting memory spaces.

Depends On D128580

[bondhugula]

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

78

Use llvm::hasSingleElement - it's O(1) instead of O(n) (since the blocks are a linked list).

187

Rephrase to

"mismatch in memory spaces between then/else ..."

?

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

bool isMemoryWrite(Operation *op, OpResult opResult,
     // implemented in terms of `bufferizesToMemoryWrite`, which does not work on
     // ops without OpOperands.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
-
-    // Compute new result types.
-    SmallVector<Type> newResultTypes;
-    for (Type type : executeRegionOp->getResultTypes()) {
-      if (auto tensorType = type.dyn_cast<TensorType>()) {
-        // TODO: Infer the result type instead of computing it.
-        newResultTypes.push_back(getMemRefType(tensorType, options));
-      } else {
-        newResultTypes.push_back(type);
-      }
-    }
+    assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
+           "only 1 block supported");
+    auto yieldOp =
+        cast<scf::YieldOp>(executeRegionOp.getRegion().front().getTerminator());
+    TypeRange newResultTypes(yieldOp.getResults());
 
     // Create new op and move over region.
     auto newOp =
         rewriter.create<scf::ExecuteRegionOp>(op->getLoc(), newResultTypes);
     newOp.getRegion().takeBody(executeRegionOp.getRegion());
 
-    // Update terminator.
-    assert(newOp.getRegion().getBlocks().size() == 1 &&
-           "only 1 block supported");
-    Block *newBlock = &newOp.getRegion().front();
-    auto yieldOp = cast<scf::YieldOp>(newBlock->getTerminator());
-    rewriter.setInsertionPoint(yieldOp);
-    SmallVector<Value> newYieldValues;
-    for (const auto &it : llvm::enumerate(yieldOp.getResults())) {
-      Value val = it.value();
-      if (val.getType().isa<TensorType>()) {
-        newYieldValues.push_back(rewriter.create<bufferization::ToMemrefOp>(
-            yieldOp.getLoc(), newResultTypes[it.index()], val));
-      } else {
-        newYieldValues.push_back(val);
-      }
-    }
-    rewriter.replaceOpWithNewOp<scf::YieldOp>(yieldOp, newYieldValues);
-
     // Update all uses of the old op.
     rewriter.setInsertionPointAfter(newOp);
     SmallVector<Value> newResults;
     for (const auto &it : llvm::enumerate(executeRegionOp->getResultTypes())) {
       if (it.value().isa<TensorType>()) {
         newResults.push_back(rewriter.create<bufferization::ToTensorOp>(
             executeRegionOp.getLoc(), newOp->getResult(it.index())));
       } else {
bool isMemoryWrite(Operation *op, OpResult opResult,
     // * conflictingWrite = %1
     //
     // For more details, check the "scf.IfOp" section of the design document.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
+    OpBuilder::InsertionGuard g(rewriter);
     auto ifOp = cast<scf::IfOp>(op);
+    auto thenYieldOp = cast<scf::YieldOp>(ifOp.thenBlock()->getTerminator());
+    auto elseYieldOp = cast<scf::YieldOp>(ifOp.elseBlock()->getTerminator());
 
-    // Compute new types of the bufferized scf.if op.
-    SmallVector<Type> newTypes;
-    for (Type returnType : ifOp->getResultTypes()) {
-      if (auto tensorType = returnType.dyn_cast<TensorType>()) {
-        // TODO: Infer the result type instead of computing it.
-        newTypes.push_back(getMemRefType(tensorType, options));
-      } else {
-        newTypes.push_back(returnType);
-      }
+    // Reconcile type mismatches between then/else branches by inserting memref
+    // casts.
+    SmallVector<Value> thenResults, elseResults;
+    bool insertedCast = false;
+    for (int64_t i = 0; i < thenYieldOp.getResults().size(); ++i) {
+      Value thenValue = thenYieldOp.getResults()[i];
+      Value elseValue = elseYieldOp.getResults()[i];
+      if (thenValue.getType() == elseValue.getType()) {
+        thenResults.push_back(thenValue);
+        elseResults.push_back(elseValue);
+        continue;
+      }
+
+      // Type mismatch between then/else yield value. Cast both to a memref type
+      // with a fully dynamic layout map.
+      auto thenMemrefType = thenValue.getType().cast<BaseMemRefType>();
+      auto elseMemrefType = elseValue.getType().cast<BaseMemRefType>();
+      if (thenMemrefType.getMemorySpaceAsInt() !=
+          elseMemrefType.getMemorySpaceAsInt())
+        return op->emitError("inconsistent memory space on then/else branches");
+      rewriter.setInsertionPoint(thenYieldOp);
+      BaseMemRefType memrefType = getMemRefTypeWithFullyDynamicLayout(
+          ifOp.getResultTypes()[i].cast<TensorType>(),
+          thenMemrefType.getMemorySpaceAsInt());
+      thenResults.push_back(rewriter.create<memref::CastOp>(
+          thenYieldOp.getLoc(), memrefType, thenValue));
+      rewriter.setInsertionPoint(elseYieldOp);
+      elseResults.push_back(rewriter.create<memref::CastOp>(
+          elseYieldOp.getLoc(), memrefType, elseValue));
+      insertedCast = true;
+
+    if (insertedCast) {
+      rewriter.setInsertionPoint(thenYieldOp);
+      rewriter.replaceOpWithNewOp<scf::YieldOp>(thenYieldOp, thenResults);
+      rewriter.setInsertionPoint(elseYieldOp);
+      rewriter.replaceOpWithNewOp<scf::YieldOp>(elseYieldOp, elseResults);
     }
 
     // Create new op.
+    rewriter.setInsertionPoint(ifOp);
+    ValueRange resultsValueRange(thenResults);
+    TypeRange newTypes(resultsValueRange);
     auto newIfOp =
         rewriter.create<scf::IfOp>(ifOp.getLoc(), newTypes, ifOp.getCondition(),
                                    /*withElseRegion=*/true);
 
-    // Remove terminators.
-    if (!newIfOp.thenBlock()->empty()) {
-      rewriter.eraseOp(newIfOp.thenBlock()->getTerminator());
-      rewriter.eraseOp(newIfOp.elseBlock()->getTerminator());
-    }
-
     // Move over then/else blocks.
     rewriter.mergeBlocks(ifOp.thenBlock(), newIfOp.thenBlock());
     rewriter.mergeBlocks(ifOp.elseBlock(), newIfOp.elseBlock());
 
-    // Update scf.yield of new then-block.
-    auto thenYieldOp = cast<scf::YieldOp>(newIfOp.thenBlock()->getTerminator());
-    rewriter.setInsertionPoint(thenYieldOp);
-    SmallVector<Value> thenYieldValues;
-    for (OpOperand &operand : thenYieldOp->getOpOperands()) {
-      if (operand.get().getType().isa<TensorType>()) {
-        ensureToMemrefOpIsValid(operand.get(),
-                                newTypes[operand.getOperandNumber()]);
-        Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
-            operand.get().getLoc(), newTypes[operand.getOperandNumber()],
-            operand.get());
-        operand.set(toMemrefOp);
-      }
-    }
-
-    // Update scf.yield of new else-block.
-    auto elseYieldOp = cast<scf::YieldOp>(newIfOp.elseBlock()->getTerminator());
-    rewriter.setInsertionPoint(elseYieldOp);
-    SmallVector<Value> elseYieldValues;
-    for (OpOperand &operand : elseYieldOp->getOpOperands()) {
-      if (operand.get().getType().isa<TensorType>()) {
-        ensureToMemrefOpIsValid(operand.get(),
-                                newTypes[operand.getOperandNumber()]);
-        Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
-            operand.get().getLoc(), newTypes[operand.getOperandNumber()],
-            operand.get());
-        operand.set(toMemrefOp);
-      }
-    }
-
     // Replace op results.
     replaceOpWithBufferizedValues(rewriter, op, newIfOp->getResults());
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
struct YieldOpInterface
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto yieldOp = cast<scf::YieldOp>(op);
     if (!isa<scf::ExecuteRegionOp, scf::IfOp, scf::ForOp, scf::WhileOp>(
             yieldOp->getParentOp()))
       return yieldOp->emitError("unsupported scf::YieldOp parent");
+
+    // TODO: Bufferize scf.yield inside scf.while/scf.for here.
+    // (Currently bufferized together with scf.while/scf.for.)
+    if (isa<scf::ForOp, scf::WhileOp>(yieldOp->getParentOp()))
+      return success();
+
+    SmallVector<Value> newResults;
+    for (const auto &it : llvm::enumerate(yieldOp.getResults())) {
+      Value value = it.value();
+      if (value.getType().isa<TensorType>()) {
+        Value buffer = getBuffer(rewriter, value, options);
+        newResults.push_back(buffer);
+      } else {
+        newResults.push_back(value);
+      }
+    }
+
+    replaceOpWithNewBufferizedOp<scf::YieldOp>(rewriter, op, newResults);
     return success();
   }
 };
 
 using tensor::ExtractSliceOp;
 
 /// Return the destinations that an ForeachThreadOp is inserting into. One per
 /// ParallelInsertSliceOp.

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-allow-return-allocs.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs allow-unknown-ops" -buffer-deallocation -canonicalize -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // CHECK-LABEL: func @buffer_not_deallocated(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
 func.func @buffer_not_deallocated(%t : tensor<?xf32>, %c : i1) -> tensor<?xf32> {
+  // CHECK: %[[m:.*]] = bufferization.to_memref %[[t]]
   // CHECK: %[[r:.*]] = scf.if %{{.*}} {
   %r = scf.if %c -> tensor<?xf32> {
     // CHECK: %[[some_op:.*]] = "test.some_op"
     // CHECK: %[[alloc:.*]] = memref.alloc(%[[some_op]])
     // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
     // CHECK-NOT: dealloc
     // CHECK: scf.yield %[[casted]]
     %sz = "test.some_op"() : () -> (index)
     %0 = bufferization.alloc_tensor(%sz) : tensor<?xf32>
     scf.yield %0 : tensor<?xf32>
   } else {
   // CHECK: } else {
-    // CHECK: %[[m:.*]] = bufferization.to_memref %[[t]]
     // CHECK: %[[cloned:.*]] = bufferization.clone %[[m]]
     // CHECK: scf.yield %[[cloned]]
     scf.yield %t : tensor<?xf32>
   }
   // CHECK: }
   // CHECK: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
   // CHECK: memref.dealloc %[[r]]
   // CHECK: return %[[r_tensor]]
   return %r : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @write_to_alloc_tensor_or_readonly_tensor(
 //  CHECK-SAME:     %[[arg0:.*]]: tensor<i32>
 func.func @write_to_alloc_tensor_or_readonly_tensor(%arg0: tensor<i32>,
                                                     %cond: i1, %val: i32)
   -> tensor<i32>
 {
-  // CHECK: %[[r:.*]] = scf.if {{.*}} {
-  // CHECK:   %[[arg0_m:.*]] = bufferization.to_memref %[[arg0]]
+  // CHECK: %[[arg0_m:.*]] = bufferization.to_memref %[[arg0]]
+  // CHECK: %[[r:.*]] = scf.if {{.*}} {
   // CHECK:   %[[clone:.*]] = bufferization.clone %[[arg0_m]]
   // CHECK:   scf.yield %[[clone]]
   // CHECK: } else {
   // CHECK:   %[[alloc:.*]] = memref.alloc
   // CHECK:   memref.store %{{.*}}, %[[alloc]]
   // CHECK:   %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK:   scf.yield %[[casted]]
   // CHECK: }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir

 }
 
 // -----
 
 // CHECK-SCF-LABEL: func @simple_scf_if(
 //  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32> {bufferization.writable = true}, %[[c:.*]]: i1, %[[pos:.*]]: index
 func.func @simple_scf_if(%t1: tensor<?xf32> {bufferization.writable = true}, %c: i1, %pos: index, %f: f32)
     -> (tensor<?xf32>, index) {
+  // CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
   // CHECK-SCF: %[[r:.*]] = scf.if %[[c]] -> (memref<?xf32, #{{.*}}>) {
   %r1, %r2 = scf.if %c -> (tensor<?xf32>, index) {
-    // CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
     // CHECK-SCF: scf.yield %[[t1_memref]]
     scf.yield %t1, %pos : tensor<?xf32>, index
   // CHECK-SCF: } else {
   } else {
     // CHECK-SCF: %[[insert:.*]] = tensor.insert %{{.*}} into %[[t1]][{{.*}}]
     // CHECK-SCF: %[[insert_memref:.*]] = bufferization.to_memref %[[insert]]
     %1 = tensor.insert %f into %t1[%pos] : tensor<?xf32>
     // CHECK-SCF: scf.yield %[[insert_memref]]
     scf.yield %1, %pos : tensor<?xf32>, index
   }
 
   // CHECK-SCF: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
   // CHECK-SCF: return %[[r_tensor]], %[[pos]]
   return %r1, %r2 : tensor<?xf32>, index
 }

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

func.func @execute_region_with_conflict(
   // CHECK: %[[alloc:.*]] = memref.alloc
   // CHECK: memref.copy %[[m1]], %[[alloc]]
   // CHECK: memref.store %{{.*}}, %[[alloc]][%{{.*}}]
   %0, %1, %2 = scf.execute_region -> (f32, tensor<?xf32>, f32) {
     %t2 = tensor.insert %f1 into %t1[%idx] : tensor<?xf32>
     scf.yield %f1, %t2, %f1 : f32, tensor<?xf32>, f32
   }
 
-  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK: %[[load:.*]] = memref.load %[[m1]]
   %3 = tensor.extract %t1[%idx] : tensor<?xf32>
 
-  // CHECK: return %{{.*}}, %[[casted]], %[[load]] : f32, memref<?xf32, #{{.*}}>, f32
+  // CHECK: return %{{.*}}, %[[alloc]], %[[load]] : f32, memref<?xf32>, f32
   return %0, %1, %3 : f32, tensor<?xf32>, f32
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_inplace(
 //  CHECK-SAME:     %[[cond:.*]]: i1, %[[t1:.*]]: memref<?xf32{{.*}}>, %[[v:.*]]: vector
 func.func @scf_if_inplace(%cond: i1,