Use SCF dialect ops for condition checking

Author: Hanumanth Hanumantharayappa

mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp

@@ -15,6 +15,7 @@
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"
 
 using namespace mlir;
@@ -296,37 +297,11 @@ struct SubViewOpInterface
       Value sizeIsNonZero = arith::CmpIOp::create(
           builder, loc, arith::CmpIPredicate::sgt, size, zero);
 
-      /*
-       * Split the current block to create the below control flow structure:
-       *
-       * ^preCondBlock:
-       *   ... // offset check already done above
-       *   %size_nonzero = arith.cmpi sgt, %size, %zero
-       *   cf.cond_br %size_nonzero, ^sizeBoundsCheckBlock, ^afterCheckBlock
-       *
-       * ^sizeBoundsCheckBlock:
-       *   %last_pos = ... // compute offset + (size-1) * stride
-       *   %last_pos_ok = ... // last position bounds check
-       *   cf.assert %last_pos_ok, "extract_slice runs out-of-bounds"
-       *   cf.br ^afterCheckBlock
-       *
-       * ^afterCheckBlock:
-       *   tensor.extract_slice ... // the original operation
-       */
-      Block *preCondBlock = builder.getBlock();
-      Block *afterCheckBlock = preCondBlock->splitBlock(subView);
-
-      // Create the block for conditional size bounds verification.
-      Block *sizeBoundsCheckBlock = builder.createBlock(
-          preCondBlock->getParent(), Region::iterator(afterCheckBlock));
-
-      // Terminate the pre-condition block with the conditional branch.
-      builder.setInsertionPointToEnd(preCondBlock);
-      cf::CondBranchOp::create(builder, loc, sizeIsNonZero,
-                               sizeBoundsCheckBlock, afterCheckBlock);
-
-      // Populate the size bounds check block with lastPos verification.
-      builder.setInsertionPointToStart(sizeBoundsCheckBlock);
+      auto ifOp = scf::IfOp::create(builder, loc, builder.getI1Type(),
+                                    sizeIsNonZero, /*withElseRegion=*/true);
+
+      // Populate the "then" region (for size > 0).
+      builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
 
       // Verify that slice does not run out-of-bounds.
       Value sizeMinusOne = arith::SubIOp::create(builder, loc, size, one);
@@ -336,12 +311,23 @@ struct SubViewOpInterface
           arith::AddIOp::create(builder, loc, offset, sizeMinusOneTimesStride);
       Value lastPosInBounds =
           generateInBoundsCheck(builder, loc, lastPos, zero, dimSize);
+
+      scf::YieldOp::create(builder, loc, lastPosInBounds);
+
+      // Populate the "else" region (for size == 0).
+      builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+      Value trueVal =
+          arith::ConstantOp::create(builder, loc, builder.getBoolAttr(true));
+      scf::YieldOp::create(builder, loc, trueVal);
+
+      builder.setInsertionPointAfter(ifOp);
+      Value finalCondition = ifOp.getResult(0);
+
       cf::AssertOp::create(
-          builder, loc, lastPosInBounds,
+          builder, loc, finalCondition,
           generateErrorMessage(op,
                                "subview runs out-of-bounds along dimension " +
                                    std::to_string(i)));
-      cf::BranchOp::create(builder, loc, afterCheckBlock);
     }
   }
 };

mlir/test/Integration/Dialect/MemRef/subview-runtime-verification.mlir

@@ -2,6 +2,7 @@
 // RUN:     -expand-strided-metadata \
 // RUN:     -lower-affine \
 // RUN:     -test-cf-assert \
+// RUN:     -convert-scf-to-cf \
 // RUN:     -convert-to-llvm | \
 // RUN: mlir-runner -e main -entry-point-result=void \
 // RUN:     -shared-libs=%mlir_runner_utils 2>&1 | \
@@ -11,6 +12,7 @@
 // RUN:     -expand-strided-metadata \
 // RUN:     -lower-affine \
 // RUN:     -test-cf-assert \
+// RUN:     -convert-scf-to-cf \
 // RUN:     -convert-to-llvm="allow-pattern-rollback=0" \
 // RUN:     -reconcile-unrealized-casts | \
 // RUN: mlir-runner -e main -entry-point-result=void \