compiler/src/iree/compiler/Dialect/Util/Transforms/Patterns.cpp - 3p/openxla/iree - Git at Google

 // Copyright 2021 The IREE Authors
 //
 // Licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

 #include "iree/compiler/Dialect/Util/Transforms/Patterns.h"

 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/IRMapping.h"

 namespace mlir::iree_compiler::IREE::Util {

 namespace {

 // Erases all |operands| that have a bit set in |elidedOperands|.
 static void eraseOperands(MutableOperandRange &operands,
                           llvm::BitVector &elidedOperands) {
   for (int i = elidedOperands.size() - 1; i >= 0; --i) {
     if (elidedOperands.test(i)) {
       operands.erase(i);
     }
   }
 }

 // Folds block arguments that are always known to have the same value at all
 // branch source sites. This is like CSE applied to block arguments.
 //
 // Example:
 //   br ^bb1(%0, %0 : index, index)
 // ^bb1(%arg0: index, %arg1: index):
 // ->
 //   br ^bb1(%0 : index)
 // ^bb1(%arg0: index):  // %arg1 remapped to %arg0
 struct FoldBlockArgumentsPattern
     : public OpInterfaceRewritePattern<CallableOpInterface> {
   using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
   LogicalResult matchAndRewrite(CallableOpInterface op,
                                 PatternRewriter &rewriter) const override {
     if (!op.getCallableRegion())
       return failure();
     auto &region = *op.getCallableRegion();
     if (region.empty() || region.hasOneBlock())
       return failure();

     // Analyze all branches in the op to compute the information we'll need to
     // analyze across branch sources.
     struct BlockSource {
       // Branch operation targeting the block.
       mutable BranchOpInterface branchOp;
       // Which successor this source represents.
       unsigned successorIndex;
       // Equivalence classes for all arguments indicating which have the same
       // value at the source. Base/non-duplicated values will be identity.
       // Example: (%a, %b, %a, %b, %c) -> (0, 2), (1, 3), (4)
       llvm::EquivalenceClasses<unsigned> duplicates;
     };
     DenseMap<Block *, SmallVector<BlockSource>> blockSourceMap;
     bool hasAnyDupes = false;
     for (auto branchOp : region.getOps<BranchOpInterface>()) {
       for (unsigned successorIndex = 0;
            successorIndex < branchOp->getNumSuccessors(); ++successorIndex) {
         auto *block = branchOp->getSuccessor(successorIndex);
         auto operands = branchOp.getSuccessorOperands(successorIndex);
         BlockSource blockSource;
         blockSource.branchOp = branchOp;
         blockSource.successorIndex = successorIndex;
         for (int i = 0; i < operands.size(); ++i) {
           blockSource.duplicates.insert(i);
           for (int j = 0; j < i; ++j) {
             if (operands[j] == operands[i]) {
               blockSource.duplicates.unionSets(i, j);
               hasAnyDupes |= true;
               break;
             }
           }
         }
         blockSourceMap[block].push_back(std::move(blockSource));
       }
     }
     if (!hasAnyDupes) {
       return failure(); // no dupes at all
     }

     rewriter.startOpModification(op);

     // Iterate over all blocks after the entry block. We can't change the entry
     // block as it is part of the function signature.
     bool didChange = false;
     for (auto &block : llvm::make_range(++region.getBlocks().begin(),
                                         region.getBlocks().end())) {
       unsigned numArgs = block.getNumArguments();
       if (numArgs == 0)
         continue;
       auto blockSources = llvm::ArrayRef(blockSourceMap[&block]);
       if (blockSources.size() == 0)
         continue;

       // Which args we'll end up erasing.
       // We need to do the actual removal after we've done the remapping below
       // as we need the values to still be live and indices consistent with the
       // analysis above.
       llvm::BitVector elidedArgs(numArgs);

       // See if each block argument is foldable across all block sources.
       // In order to fold we need each source to share some duplicates but note
       // that the sources may not have identical sets.
       llvm::BitVector sameValues(numArgs);   // reused
       llvm::BitVector sourceValues(numArgs); // reused
       for (unsigned argIndex = 0; argIndex < numArgs; ++argIndex) {
         // Each bit represents an argument that duplicates the arg at argIndex.
         // We walk all the sources and AND their masks together to get the safe
         // set of duplicate operands.
         // Example for %0: (%a, %b, %a) -> b001
         // Example for %1: (%a, %b, %a) -> b000
         sameValues.set(); // note reused
         for (auto &blockSource : blockSources) {
           sourceValues.reset();
           for (auto mit = blockSource.duplicates.findLeader(argIndex);
                mit != blockSource.duplicates.member_end(); ++mit) {
             sourceValues.set(*mit);
           }
           sameValues &= sourceValues;
         }
         if (sameValues.none()) {
           continue; // arg unused/not duplicated
         }

         // Remove the base argument from the set so we don't erase it and can
         // point all duplicate args at it.
         int baseArgIndex = sameValues.find_first();
         sameValues.reset(baseArgIndex);
         elidedArgs |= sameValues;

         // Replace all of the subsequent duplicate arguments with the first.
         auto baseArg = block.getArgument(baseArgIndex);
         for (unsigned dupeIndex : sameValues.set_bits()) {
           rewriter.replaceAllUsesWith(block.getArgument(dupeIndex), baseArg);
         }
       }

       // Erase all the block arguments we've deduplicated.
       if (elidedArgs.any()) {
         for (auto &blockSource : blockSources) {
           auto successorOperands = blockSource.branchOp.getSuccessorOperands(
               blockSource.successorIndex);
           auto operands = successorOperands.slice(
               successorOperands.getProducedOperandCount(),
               successorOperands.size());
           rewriter.modifyOpInPlace(blockSource.branchOp, [&]() {
             eraseOperands(operands, elidedArgs);
           });
         }
         block.eraseArguments(elidedArgs);
         didChange |= !elidedArgs.none();
       }
     }

     if (didChange) {
       rewriter.finalizeOpModification(op);
       return success();
     } else {
       rewriter.cancelOpModification(op);
       return failure();
     }
   }
 };

 // Finds branch arguments that come from dominating branch source operands on
 // all incoming edges and elides the arguments.
 //
 // Example:
 //  func.func @foo(%arg0: index) {
 //    br ^bb1(%arg0 : index)
 //  ^bb1(%0: index):
 // ->
 //  func.func @foo(%arg0: index) {
 //    br ^bb1
 //  ^bb1:  // %0 remapped to %arg0
 struct ElideBranchOperandsPattern
     : public OpInterfaceRewritePattern<CallableOpInterface> {
   using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
   LogicalResult matchAndRewrite(CallableOpInterface op,
                                 PatternRewriter &rewriter) const override {
     if (!op.getCallableRegion())
       return failure();
     auto &region = *op.getCallableRegion();
     if (region.empty())
       return failure();
     DominanceInfo dominance(op);

     // Analyze all branches to build a map of blocks to their sources.
     struct BlockSource {
       // Branch operation targeting the block.
       mutable BranchOpInterface branchOp;
       // Which successor this source represents.
       unsigned successorIndex;
     };
     DenseMap<Block *, SmallVector<BlockSource>> blockSourceMap;
     for (auto branchOp : region.getOps<BranchOpInterface>()) {
       for (unsigned successorIndex = 0;
            successorIndex < branchOp->getNumSuccessors(); ++successorIndex) {
         auto *block = branchOp->getSuccessor(successorIndex);
         auto operands = branchOp.getSuccessorOperands(successorIndex);
         BlockSource blockSource;
         blockSource.branchOp = branchOp;
         blockSource.successorIndex = successorIndex;
         blockSourceMap[block].push_back(blockSource);
       }
     }

     rewriter.startOpModification(op);

     // Iterate over all blocks after the entry block. We can't change the entry
     // block as it is part of the function signature.
     bool didChange = false;
     for (auto &block : llvm::make_range(++region.getBlocks().begin(),
                                         region.getBlocks().end())) {
       unsigned numArgs = block.getNumArguments();
       if (numArgs == 0)
         continue;
       auto blockSources = llvm::ArrayRef(blockSourceMap[&block]);
       if (blockSources.size() == 0)
         continue;

       // Which args we'll end up erasing.
       // We need to do the actual removal after we've done the remapping below
       // as we need the values to still be live and indices consistent with the
       // analysis above.
       llvm::BitVector elidedArgs(numArgs);

       for (unsigned argIndex = 0; argIndex < numArgs; ++argIndex) {
         // Find the uniform value passed for the operand of all branches.
         Value uniformValue = nullptr;
         for (auto &blockSource : blockSources) {
           auto operands = blockSource.branchOp.getSuccessorOperands(
               blockSource.successorIndex);
           auto operand = operands[argIndex];
           if (!uniformValue) {
             // First usage.
             uniformValue = operand;
             continue;
           } else if (uniformValue == operand) {
             // Operands match between all previous and the current source.
             continue;
           }

           // Operand for this source differs from previous. This is either
           // because it's non-uniform _or_ that it's a cycle.
           if (auto sourceArg = llvm::dyn_cast<BlockArgument>(operand)) {
             // Operand comes from a block argument. If that is the block
             // argument we are analyzing it means there's a cycle (%0 -> %0) and
             // we can ignore it for the purposes of this analysis.
             if (sourceArg.getOwner() == &block &&
                 sourceArg.getArgNumber() == argIndex) {
               continue;
             }
           }

           // Non-uniform; can't elide.
           uniformValue = nullptr;
           break;
         }
         if (!uniformValue)
           continue;

         // See if the uniform value dominates this block; if so we can use it.
         if (!uniformValue.getDefiningOp() ||
             dominance.dominates(uniformValue.getDefiningOp()->getBlock(),
                                 &block)) {
           rewriter.replaceAllUsesWith(block.getArgument(argIndex),
                                       uniformValue);
           elidedArgs.set(argIndex);
         }
       }
       if (elidedArgs.none())
         continue;

       // Erase all the block arguments we remapped.
       for (auto &blockSource : blockSources) {
         auto successorOperands = blockSource.branchOp.getSuccessorOperands(
             blockSource.successorIndex);
         auto operands =
             successorOperands.slice(successorOperands.getProducedOperandCount(),
                                     successorOperands.size());
         rewriter.modifyOpInPlace(blockSource.branchOp, [&]() {
           eraseOperands(operands, elidedArgs);
         });
       }
       block.eraseArguments(elidedArgs);
       didChange |= !elidedArgs.none();
     }

     if (didChange) {
       rewriter.finalizeOpModification(op);
       return success();
     } else {
       rewriter.cancelOpModification(op);
       return failure();
     }
   }
 };

 // Converts an scf.index_switch with a single case into an scf.if.
 //
 // Example:
 //  scf.index_switch %case : i32
 //  case 0 {
 //    %foo = ...
 //    scf.yield %foo : i32
 //  }
 //  default {
 //    %default = ...
 //    scf.yield %default : i32
 //  }
 // ->
 //  %case_0 = arith.cmpi eq, %case, %c0 : index
 //  scf.if %case_0 -> i32 {
 //    %foo = ...
 //    scf.yield %foo : i32
 //  } else {
 //    %default = ...
 //    scf.yield %default : i32
 //  }
 struct IndexSwitchToIfPattern : public OpRewritePattern<scf::IndexSwitchOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(scf::IndexSwitchOp switchOp,
                                 PatternRewriter &rewriter) const override {
     if (switchOp.getNumCases() != 1)
       return failure();
     Value caseValue = rewriter.create<arith::ConstantIndexOp>(
         switchOp.getLoc(), switchOp.getCases().front());
     Value isCaseValue = rewriter.createOrFold<arith::CmpIOp>(
         switchOp.getLoc(), arith::CmpIPredicate::eq, switchOp.getArg(),
         caseValue);
     auto ifOp = rewriter.create<scf::IfOp>(
         switchOp.getLoc(), switchOp.getResultTypes(), isCaseValue);
     rewriter.inlineRegionBefore(switchOp.getCaseRegions().front(),
                                 ifOp.getThenRegion(),
                                 ifOp.getThenRegion().begin());
     rewriter.inlineRegionBefore(switchOp.getDefaultRegion(),
                                 ifOp.getElseRegion(),
                                 ifOp.getElseRegion().begin());
     rewriter.replaceOp(switchOp, ifOp);
     return success();
   }
 };

 // Combines adjacent and compatible scf.index_switch ops into one.
 // We sink any ops from the first switch into the second and erase the first.
 // Note that since the second may implicitly capture the results of the first
 // we have to handle mapping the scf.yield results to their new local values in
 // the second.
 //
 // Example:
 //  scf.index_switch %case
 //  case 0 {
 //    foo
 //    scf.yield
 //  }
 //  default {
 //    foo_default
 //    scf.yield
 //  }
 //  scf.index_switch %case
 //  case 0 {
 //    bar
 //    scf.yield
 //  }
 //  default {
 //    bar_default
 //    scf.yield
 //  }
 // ->
 //  scf.index_switch %case
 //  case 0 {
 //    foo
 //    bar
 //    scf.yield
 //  }
 //  default {
 //    foo_default
 //    bar_default
 //    scf.yield
 //  }
 struct MergeIndexSwitchPattern : public OpRewritePattern<scf::IndexSwitchOp> {
   MergeIndexSwitchPattern(MLIRContext *context)
       : OpRewritePattern(context, 1000) {}
   LogicalResult matchAndRewrite(scf::IndexSwitchOp nextOp,
                                 PatternRewriter &rewriter) const override {
     // Inspect the previous op to see if it's also a switch.
     auto prevOp = dyn_cast_or_null<scf::IndexSwitchOp>(nextOp->getPrevNode());
     if (!prevOp)
       return failure();

     // Require that the cases line up exactly. There's probably some merging
     // we could do in other cases but it'd be best to leave other patterns to
     // hoist/CSE cases/etc instead.
     if (prevOp.getNumCases() != nextOp.getNumCases())
       return rewriter.notifyMatchFailure(nextOp, "number of cases differ");
     if (!llvm::equal(prevOp.getCases(), nextOp.getCases()))
       return rewriter.notifyMatchFailure(nextOp, "case values differ");

     // Create a new switch to replace nextOp that contains the same cases but
     // combined results from both ops.
     SmallVector<Type> newResultTypes;
     llvm::append_range(newResultTypes, prevOp.getResultTypes());
     llvm::append_range(newResultTypes, nextOp.getResultTypes());
     auto newOp = rewriter.create<scf::IndexSwitchOp>(
         rewriter.getFusedLoc({prevOp.getLoc(), nextOp.getLoc()}),
         newResultTypes, prevOp.getArg(), prevOp.getCases(),
         prevOp.getNumCases());
     SmallVector<std::pair<Value, Value>> resultReplacements;
     for (auto [oldResult, newResult] :
          llvm::zip_equal(prevOp.getResults(),
                          newOp.getResults().slice(0, prevOp.getNumResults()))) {
       resultReplacements.push_back(std::make_pair(oldResult, newResult));
     }
     for (auto [oldResult, newResult] :
          llvm::zip_equal(nextOp.getResults(),
                          newOp.getResults().slice(prevOp.getNumResults(),
                                                   nextOp.getNumResults()))) {
       resultReplacements.push_back(std::make_pair(oldResult, newResult));
     }

     // NOTE: the results of prevOp may be implicitly captured by nextOp and we
     // need to build a mapping of the prevOp results to their cloned values in
     // nextOp once merged.

     auto cloneRegions = [&](Region &regionA, Region &regionB, Region &target) {
       SmallVector<Value> yieldValues;
       IRMapping localMapping;
       auto targetBuilder = OpBuilder::atBlockBegin(&target.emplaceBlock());

       // Clone regionA into target and map any results from prevOp to the cloned
       // values for the particular case.
       auto yieldA = *regionA.getOps<scf::YieldOp>().begin();
       for (auto &op : regionA.getOps()) {
         if (op.hasTrait<OpTrait::IsTerminator>())
           continue;
         // Clone each op and map its original value to the new local value.
         targetBuilder.clone(op, localMapping);
       }
       for (auto [yieldValue, resultValue] :
            llvm::zip_equal(yieldA.getOperands(), prevOp.getResults())) {
         // Track the new local values yielded by the region.
         auto newValue = localMapping.lookupOrDefault(yieldValue);
         yieldValues.push_back(newValue);
         localMapping.map(resultValue, newValue);
       }

       // Clone regionB into target.
       auto yieldB = *regionB.getOps<scf::YieldOp>().begin();
       for (auto &op : regionB.getOps()) {
         if (op.hasTrait<OpTrait::IsTerminator>())
           continue;
         // Clone each op and map its original value to the new local value.
         targetBuilder.clone(op, localMapping);
       }
       for (auto yieldValue : yieldB.getOperands()) {
         // Track the new local values yielded by the region, ensuring we check
         // what the first region may have produced and been captured implicitly.
         yieldValues.push_back(localMapping.lookupOrNull(yieldValue));
       }

       // Add the merged yield containing results from both regions.
       targetBuilder.create<scf::YieldOp>(
           targetBuilder.getFusedLoc(yieldA.getLoc(), yieldB.getLoc()),
           yieldValues);
     };

     // Merge regions from both prevOp and nextOp into the newOp.
     for (auto [prevRegion, nextRegion, newRegion] :
          llvm::zip_equal(prevOp.getCaseRegions(), nextOp.getCaseRegions(),
                          newOp.getCaseRegions())) {
       cloneRegions(prevRegion, nextRegion, newRegion);
     }
     cloneRegions(prevOp.getDefaultRegion(), nextOp.getDefaultRegion(),
                  newOp.getDefaultRegion());

     // Update uses of the old results from both ops to the new ones.
     for (auto [oldResult, newResult] : resultReplacements) {
       rewriter.replaceAllUsesWith(oldResult, newResult);
     }
     rewriter.eraseOp(prevOp);
     rewriter.eraseOp(nextOp);

     return success();
   }
 };

 } // namespace

 void populateCommonPatterns(MLIRContext *context, RewritePatternSet &patterns) {
   context->getOrLoadDialect<IREE::Util::UtilDialect>()
       ->getCanonicalizationPatterns(patterns);

   patterns.insert<FoldBlockArgumentsPattern, ElideBranchOperandsPattern>(
       context);

   patterns.insert<IndexSwitchToIfPattern, MergeIndexSwitchPattern>(context);
 }

 } // namespace mlir::iree_compiler::IREE::Util
	// Copyright 2021 The IREE Authors
	//
	// Licensed under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

	#include "iree/compiler/Dialect/Util/Transforms/Patterns.h"

	#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
	#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
	#include "llvm/ADT/BitVector.h"
	#include "llvm/ADT/EquivalenceClasses.h"
	#include "mlir/Dialect/SCF/IR/SCF.h"
	#include "mlir/IR/BuiltinOps.h"
	#include "mlir/IR/Dominance.h"
	#include "mlir/IR/IRMapping.h"

	namespace mlir::iree_compiler::IREE::Util {

	namespace {

	// Erases all \|operands\| that have a bit set in \|elidedOperands\|.
	static void eraseOperands(MutableOperandRange &operands,
	llvm::BitVector &elidedOperands) {
	for (int i = elidedOperands.size() - 1; i >= 0; --i) {
	if (elidedOperands.test(i)) {
	operands.erase(i);
	}
	}
	}

	// Folds block arguments that are always known to have the same value at all
	// branch source sites. This is like CSE applied to block arguments.
	//
	// Example:
	// br ^bb1(%0, %0 : index, index)
	// ^bb1(%arg0: index, %arg1: index):
	// ->
	// br ^bb1(%0 : index)
	// ^bb1(%arg0: index): // %arg1 remapped to %arg0
	struct FoldBlockArgumentsPattern
	: public OpInterfaceRewritePattern<CallableOpInterface> {
	using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
	LogicalResult matchAndRewrite(CallableOpInterface op,
	PatternRewriter &rewriter) const override {
	if (!op.getCallableRegion())
	return failure();
	auto &region = *op.getCallableRegion();
	if (region.empty() \|\| region.hasOneBlock())
	return failure();

	// Analyze all branches in the op to compute the information we'll need to
	// analyze across branch sources.
	struct BlockSource {
	// Branch operation targeting the block.
	mutable BranchOpInterface branchOp;
	// Which successor this source represents.
	unsigned successorIndex;
	// Equivalence classes for all arguments indicating which have the same
	// value at the source. Base/non-duplicated values will be identity.
	// Example: (%a, %b, %a, %b, %c) -> (0, 2), (1, 3), (4)
	llvm::EquivalenceClasses<unsigned> duplicates;
	};
	DenseMap<Block *, SmallVector<BlockSource>> blockSourceMap;
	bool hasAnyDupes = false;
	for (auto branchOp : region.getOps<BranchOpInterface>()) {
	for (unsigned successorIndex = 0;
	successorIndex < branchOp->getNumSuccessors(); ++successorIndex) {
	auto *block = branchOp->getSuccessor(successorIndex);
	auto operands = branchOp.getSuccessorOperands(successorIndex);
	BlockSource blockSource;
	blockSource.branchOp = branchOp;
	blockSource.successorIndex = successorIndex;
	for (int i = 0; i < operands.size(); ++i) {
	blockSource.duplicates.insert(i);
	for (int j = 0; j < i; ++j) {
	if (operands[j] == operands[i]) {
	blockSource.duplicates.unionSets(i, j);
	hasAnyDupes \|= true;
	break;
	}
	}
	}
	blockSourceMap[block].push_back(std::move(blockSource));
	}
	}
	if (!hasAnyDupes) {
	return failure(); // no dupes at all
	}

	rewriter.startOpModification(op);

	// Iterate over all blocks after the entry block. We can't change the entry
	// block as it is part of the function signature.
	bool didChange = false;
	for (auto &block : llvm::make_range(++region.getBlocks().begin(),
	region.getBlocks().end())) {
	unsigned numArgs = block.getNumArguments();
	if (numArgs == 0)
	continue;
	auto blockSources = llvm::ArrayRef(blockSourceMap[&block]);
	if (blockSources.size() == 0)
	continue;

	// Which args we'll end up erasing.
	// We need to do the actual removal after we've done the remapping below
	// as we need the values to still be live and indices consistent with the
	// analysis above.
	llvm::BitVector elidedArgs(numArgs);

	// See if each block argument is foldable across all block sources.
	// In order to fold we need each source to share some duplicates but note
	// that the sources may not have identical sets.
	llvm::BitVector sameValues(numArgs); // reused
	llvm::BitVector sourceValues(numArgs); // reused
	for (unsigned argIndex = 0; argIndex < numArgs; ++argIndex) {
	// Each bit represents an argument that duplicates the arg at argIndex.
	// We walk all the sources and AND their masks together to get the safe
	// set of duplicate operands.
	// Example for %0: (%a, %b, %a) -> b001
	// Example for %1: (%a, %b, %a) -> b000
	sameValues.set(); // note reused
	for (auto &blockSource : blockSources) {
	sourceValues.reset();
	for (auto mit = blockSource.duplicates.findLeader(argIndex);
	mit != blockSource.duplicates.member_end(); ++mit) {
	sourceValues.set(*mit);
	}
	sameValues &= sourceValues;
	}
	if (sameValues.none()) {
	continue; // arg unused/not duplicated
	}

	// Remove the base argument from the set so we don't erase it and can
	// point all duplicate args at it.
	int baseArgIndex = sameValues.find_first();
	sameValues.reset(baseArgIndex);
	elidedArgs \|= sameValues;

	// Replace all of the subsequent duplicate arguments with the first.
	auto baseArg = block.getArgument(baseArgIndex);
	for (unsigned dupeIndex : sameValues.set_bits()) {
	rewriter.replaceAllUsesWith(block.getArgument(dupeIndex), baseArg);
	}
	}

	// Erase all the block arguments we've deduplicated.
	if (elidedArgs.any()) {
	for (auto &blockSource : blockSources) {
	auto successorOperands = blockSource.branchOp.getSuccessorOperands(
	blockSource.successorIndex);
	auto operands = successorOperands.slice(
	successorOperands.getProducedOperandCount(),
	successorOperands.size());
	rewriter.modifyOpInPlace(blockSource.branchOp, [&]() {
	eraseOperands(operands, elidedArgs);
	});
	}
	block.eraseArguments(elidedArgs);
	didChange \|= !elidedArgs.none();
	}
	}

	if (didChange) {
	rewriter.finalizeOpModification(op);
	return success();
	} else {
	rewriter.cancelOpModification(op);
	return failure();
	}
	}
	};

	// Finds branch arguments that come from dominating branch source operands on
	// all incoming edges and elides the arguments.
	//
	// Example:
	// func.func @foo(%arg0: index) {
	// br ^bb1(%arg0 : index)
	// ^bb1(%0: index):
	// ->
	// func.func @foo(%arg0: index) {
	// br ^bb1
	// ^bb1: // %0 remapped to %arg0
	struct ElideBranchOperandsPattern
	: public OpInterfaceRewritePattern<CallableOpInterface> {
	using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
	LogicalResult matchAndRewrite(CallableOpInterface op,
	PatternRewriter &rewriter) const override {
	if (!op.getCallableRegion())
	return failure();
	auto &region = *op.getCallableRegion();
	if (region.empty())
	return failure();
	DominanceInfo dominance(op);

	// Analyze all branches to build a map of blocks to their sources.
	struct BlockSource {
	// Branch operation targeting the block.
	mutable BranchOpInterface branchOp;
	// Which successor this source represents.
	unsigned successorIndex;
	};
	DenseMap<Block *, SmallVector<BlockSource>> blockSourceMap;
	for (auto branchOp : region.getOps<BranchOpInterface>()) {
	for (unsigned successorIndex = 0;
	successorIndex < branchOp->getNumSuccessors(); ++successorIndex) {
	auto *block = branchOp->getSuccessor(successorIndex);
	auto operands = branchOp.getSuccessorOperands(successorIndex);
	BlockSource blockSource;
	blockSource.branchOp = branchOp;
	blockSource.successorIndex = successorIndex;
	blockSourceMap[block].push_back(blockSource);
	}
	}

	rewriter.startOpModification(op);

	// Iterate over all blocks after the entry block. We can't change the entry
	// block as it is part of the function signature.
	bool didChange = false;
	for (auto &block : llvm::make_range(++region.getBlocks().begin(),
	region.getBlocks().end())) {
	unsigned numArgs = block.getNumArguments();
	if (numArgs == 0)
	continue;
	auto blockSources = llvm::ArrayRef(blockSourceMap[&block]);
	if (blockSources.size() == 0)
	continue;

	// Which args we'll end up erasing.
	// We need to do the actual removal after we've done the remapping below
	// as we need the values to still be live and indices consistent with the
	// analysis above.
	llvm::BitVector elidedArgs(numArgs);

	for (unsigned argIndex = 0; argIndex < numArgs; ++argIndex) {
	// Find the uniform value passed for the operand of all branches.
	Value uniformValue = nullptr;
	for (auto &blockSource : blockSources) {
	auto operands = blockSource.branchOp.getSuccessorOperands(
	blockSource.successorIndex);
	auto operand = operands[argIndex];
	if (!uniformValue) {
	// First usage.
	uniformValue = operand;
	continue;
	} else if (uniformValue == operand) {
	// Operands match between all previous and the current source.
	continue;
	}

	// Operand for this source differs from previous. This is either
	// because it's non-uniform _or_ that it's a cycle.
	if (auto sourceArg = llvm::dyn_cast<BlockArgument>(operand)) {
	// Operand comes from a block argument. If that is the block
	// argument we are analyzing it means there's a cycle (%0 -> %0) and
	// we can ignore it for the purposes of this analysis.
	if (sourceArg.getOwner() == &block &&
	sourceArg.getArgNumber() == argIndex) {
	continue;
	}
	}

	// Non-uniform; can't elide.
	uniformValue = nullptr;
	break;
	}
	if (!uniformValue)
	continue;

	// See if the uniform value dominates this block; if so we can use it.
	if (!uniformValue.getDefiningOp() \|\|
	dominance.dominates(uniformValue.getDefiningOp()->getBlock(),
	&block)) {
	rewriter.replaceAllUsesWith(block.getArgument(argIndex),
	uniformValue);
	elidedArgs.set(argIndex);
	}
	}
	if (elidedArgs.none())
	continue;

	// Erase all the block arguments we remapped.
	for (auto &blockSource : blockSources) {
	auto successorOperands = blockSource.branchOp.getSuccessorOperands(
	blockSource.successorIndex);
	auto operands =
	successorOperands.slice(successorOperands.getProducedOperandCount(),
	successorOperands.size());
	rewriter.modifyOpInPlace(blockSource.branchOp, [&]() {
	eraseOperands(operands, elidedArgs);
	});
	}
	block.eraseArguments(elidedArgs);
	didChange \|= !elidedArgs.none();
	}

	if (didChange) {
	rewriter.finalizeOpModification(op);
	return success();
	} else {
	rewriter.cancelOpModification(op);
	return failure();
	}
	}
	};

	// Converts an scf.index_switch with a single case into an scf.if.
	//
	// Example:
	// scf.index_switch %case : i32
	// case 0 {
	// %foo = ...
	// scf.yield %foo : i32
	// }
	// default {
	// %default = ...
	// scf.yield %default : i32
	// }
	// ->
	// %case_0 = arith.cmpi eq, %case, %c0 : index
	// scf.if %case_0 -> i32 {
	// %foo = ...
	// scf.yield %foo : i32
	// } else {
	// %default = ...
	// scf.yield %default : i32
	// }
	struct IndexSwitchToIfPattern : public OpRewritePattern<scf::IndexSwitchOp> {
	using OpRewritePattern::OpRewritePattern;
	LogicalResult matchAndRewrite(scf::IndexSwitchOp switchOp,
	PatternRewriter &rewriter) const override {
	if (switchOp.getNumCases() != 1)
	return failure();
	Value caseValue = rewriter.create<arith::ConstantIndexOp>(
	switchOp.getLoc(), switchOp.getCases().front());
	Value isCaseValue = rewriter.createOrFold<arith::CmpIOp>(
	switchOp.getLoc(), arith::CmpIPredicate::eq, switchOp.getArg(),
	caseValue);
	auto ifOp = rewriter.create<scf::IfOp>(
	switchOp.getLoc(), switchOp.getResultTypes(), isCaseValue);
	rewriter.inlineRegionBefore(switchOp.getCaseRegions().front(),
	ifOp.getThenRegion(),
	ifOp.getThenRegion().begin());
	rewriter.inlineRegionBefore(switchOp.getDefaultRegion(),
	ifOp.getElseRegion(),
	ifOp.getElseRegion().begin());
	rewriter.replaceOp(switchOp, ifOp);
	return success();
	}
	};

	// Combines adjacent and compatible scf.index_switch ops into one.
	// We sink any ops from the first switch into the second and erase the first.
	// Note that since the second may implicitly capture the results of the first
	// we have to handle mapping the scf.yield results to their new local values in
	// the second.
	//
	// Example:
	// scf.index_switch %case
	// case 0 {
	// foo
	// scf.yield
	// }
	// default {
	// foo_default
	// scf.yield
	// }
	// scf.index_switch %case
	// case 0 {
	// bar
	// scf.yield
	// }
	// default {
	// bar_default
	// scf.yield
	// }
	// ->
	// scf.index_switch %case
	// case 0 {
	// foo
	// bar
	// scf.yield
	// }
	// default {
	// foo_default
	// bar_default
	// scf.yield
	// }
	struct MergeIndexSwitchPattern : public OpRewritePattern<scf::IndexSwitchOp> {
	MergeIndexSwitchPattern(MLIRContext *context)
	: OpRewritePattern(context, 1000) {}
	LogicalResult matchAndRewrite(scf::IndexSwitchOp nextOp,
	PatternRewriter &rewriter) const override {
	// Inspect the previous op to see if it's also a switch.
	auto prevOp = dyn_cast_or_null<scf::IndexSwitchOp>(nextOp->getPrevNode());
	if (!prevOp)
	return failure();

	// Require that the cases line up exactly. There's probably some merging
	// we could do in other cases but it'd be best to leave other patterns to
	// hoist/CSE cases/etc instead.
	if (prevOp.getNumCases() != nextOp.getNumCases())
	return rewriter.notifyMatchFailure(nextOp, "number of cases differ");
	if (!llvm::equal(prevOp.getCases(), nextOp.getCases()))
	return rewriter.notifyMatchFailure(nextOp, "case values differ");

	// Create a new switch to replace nextOp that contains the same cases but
	// combined results from both ops.
	SmallVector<Type> newResultTypes;
	llvm::append_range(newResultTypes, prevOp.getResultTypes());
	llvm::append_range(newResultTypes, nextOp.getResultTypes());
	auto newOp = rewriter.create<scf::IndexSwitchOp>(
	rewriter.getFusedLoc({prevOp.getLoc(), nextOp.getLoc()}),
	newResultTypes, prevOp.getArg(), prevOp.getCases(),
	prevOp.getNumCases());
	SmallVector<std::pair<Value, Value>> resultReplacements;
	for (auto [oldResult, newResult] :
	llvm::zip_equal(prevOp.getResults(),
	newOp.getResults().slice(0, prevOp.getNumResults()))) {
	resultReplacements.push_back(std::make_pair(oldResult, newResult));
	}
	for (auto [oldResult, newResult] :
	llvm::zip_equal(nextOp.getResults(),
	newOp.getResults().slice(prevOp.getNumResults(),
	nextOp.getNumResults()))) {
	resultReplacements.push_back(std::make_pair(oldResult, newResult));
	}

	// NOTE: the results of prevOp may be implicitly captured by nextOp and we
	// need to build a mapping of the prevOp results to their cloned values in
	// nextOp once merged.

	auto cloneRegions = [&](Region &regionA, Region &regionB, Region &target) {
	SmallVector<Value> yieldValues;
	IRMapping localMapping;
	auto targetBuilder = OpBuilder::atBlockBegin(&target.emplaceBlock());

	// Clone regionA into target and map any results from prevOp to the cloned
	// values for the particular case.
	auto yieldA = *regionA.getOps<scf::YieldOp>().begin();
	for (auto &op : regionA.getOps()) {
	if (op.hasTrait<OpTrait::IsTerminator>())
	continue;
	// Clone each op and map its original value to the new local value.
	targetBuilder.clone(op, localMapping);
	}
	for (auto [yieldValue, resultValue] :
	llvm::zip_equal(yieldA.getOperands(), prevOp.getResults())) {
	// Track the new local values yielded by the region.
	auto newValue = localMapping.lookupOrDefault(yieldValue);
	yieldValues.push_back(newValue);
	localMapping.map(resultValue, newValue);
	}

	// Clone regionB into target.
	auto yieldB = *regionB.getOps<scf::YieldOp>().begin();
	for (auto &op : regionB.getOps()) {
	if (op.hasTrait<OpTrait::IsTerminator>())
	continue;
	// Clone each op and map its original value to the new local value.
	targetBuilder.clone(op, localMapping);
	}
	for (auto yieldValue : yieldB.getOperands()) {
	// Track the new local values yielded by the region, ensuring we check
	// what the first region may have produced and been captured implicitly.
	yieldValues.push_back(localMapping.lookupOrNull(yieldValue));
	}

	// Add the merged yield containing results from both regions.
	targetBuilder.create<scf::YieldOp>(
	targetBuilder.getFusedLoc(yieldA.getLoc(), yieldB.getLoc()),
	yieldValues);
	};

	// Merge regions from both prevOp and nextOp into the newOp.
	for (auto [prevRegion, nextRegion, newRegion] :
	llvm::zip_equal(prevOp.getCaseRegions(), nextOp.getCaseRegions(),
	newOp.getCaseRegions())) {
	cloneRegions(prevRegion, nextRegion, newRegion);
	}
	cloneRegions(prevOp.getDefaultRegion(), nextOp.getDefaultRegion(),
	newOp.getDefaultRegion());

	// Update uses of the old results from both ops to the new ones.
	for (auto [oldResult, newResult] : resultReplacements) {
	rewriter.replaceAllUsesWith(oldResult, newResult);
	}
	rewriter.eraseOp(prevOp);
	rewriter.eraseOp(nextOp);

	return success();
	}
	};

	} // namespace

	void populateCommonPatterns(MLIRContext *context, RewritePatternSet &patterns) {
	context->getOrLoadDialect<IREE::Util::UtilDialect>()
	->getCanonicalizationPatterns(patterns);

	patterns.insert<FoldBlockArgumentsPattern, ElideBranchOperandsPattern>(
	context);

	patterns.insert<IndexSwitchToIfPattern, MergeIndexSwitchPattern>(context);
	}

	} // namespace mlir::iree_compiler::IREE::Util