compiler/src/iree/compiler/DispatchCreation/FormScalarDispatches.cpp - 3p/openxla/iree - Git at Google

 // Copyright 2023 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/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Debug.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Iterators.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"

 #define DEBUG_TYPE "iree-dispatch-creation-form-scalar-dispatches"

 namespace mlir::iree_compiler::DispatchCreation {

 #define GEN_PASS_DEF_FORMSCALARDISPATCHESPASS
 #include "iree/compiler/DispatchCreation/Passes.h.inc"

 namespace {

 /// Pass declaration.
 struct FormScalarDispatchesPass final
     : public impl::FormScalarDispatchesPassBase<FormScalarDispatchesPass> {
   void runOnOperation() override;
 };
 } // namespace

 /// Return true if type represents a value less than `n` elements.
 static bool isScalarOrTensorOfLinearSizeN(int n, Type type) {
   if (type.isIntOrIndexOrFloat()) {
     return true;
   }
   if (auto tensorType = dyn_cast<RankedTensorType>(type)) {
     if (!tensorType.hasStaticShape()) {
       return false;
     }
     return tensorType.getNumElements() <= n;
   }
   return false;
 }

 /// Return `true` for operations that are to be treated as compute operations.
 static bool isComputeOperation(Operation *op) {
   MLIRContext *context = op->getContext();
   if (op->getDialect() == context->getLoadedDialect<linalg::LinalgDialect>()) {
     return true;
   }
   if (op->getDialect() == context->getLoadedDialect<tensor::TensorDialect>()) {
     return !isa<tensor::CastOp, tensor::CollapseShapeOp, tensor::EmptyOp,
                 tensor::ExpandShapeOp, tensor::PackOp, tensor::UnPackOp>(op);
   }
   return false;
 }

 /// Return `true` if the workload of this operation is less than `n`.
 static bool isOperationWorkloadLessThanSizeN(int n, Operation *candidateOp) {
   return llvm::all_of(candidateOp->getOperands(),
                       [&](Value v) {
                         return isScalarOrTensorOfLinearSizeN(n, v.getType());
                       }) &&
          llvm::all_of(candidateOp->getResultTypes(), [&](Type t) {
            return isScalarOrTensorOfLinearSizeN(n, t);
          });
 }

 /// Return `true` is the operation is to be treated as a scalar operation
 /// and moved into a scalar dispatch (not necessarily as the root of the
 /// dispatch).
 static bool isScalarOperation(int workload, Operation *op) {
   // 1. Ignore most operations. Only look for a whitelist set of operations.
   if (!isComputeOperation(op)) {
     return false;
   }

   // 2. Check that the workload of the operation is less then the limit
   if (!isOperationWorkloadLessThanSizeN(workload, op)) {
     return false;
   }

   // 3. Do not move operations that are cloned into the dispatch region.
   // TODO: This might prevent moving all scalar operations into dispatch
   // resulting in artifical splits. Revisit after more examples.
   return !IREE::Flow::isClonableIntoDispatchOp(op);
 }

 /// Given a `rootOp` return a DAG of the program that represents
 /// operations that can be moved into a scalar dispatch with the `rootOp`
 /// as the root of the DAG.
 llvm::SetVector<Operation *> computeSliceToMoveIntoDispatch(
     int workload, Operation *rootOp,
     const llvm::DenseMap<Operation *, Operation *> &opToRootMap) {
   BackwardSliceOptions options;
   options.filter = [&](Operation *currentOp) {
     assert(currentOp && "current op is null");
     if (opToRootMap.count(currentOp)) {
       return false;
     }
     // Operations needs to be in the same block as `rootOp`.
     if (currentOp->getBlock() != rootOp->getBlock()) {
       return false;
     }

     if (!isScalarOperation(workload, currentOp)) {
       return false;
     }

     // All its uses must be in the `opToRootMap`, i.e. they are either
     // in the current dispatches, or those already formed.
     return llvm::all_of(currentOp->getUsers(), [&](Operation *user) {
       return opToRootMap.count(user);
     });
   };
   options.omitBlockArguments = true;
   llvm::SetVector<Operation *> slice;
   getBackwardSlice(rootOp, &slice, options);
   return slice;
 }

 /// Return `true` if the op is to be treated as a root of a scalar dispatch.
 static bool isSliceRoot(int workload, Operation *op) {
   return !op->getParentOfType<IREE::Flow::DispatchRegionOp>() &&
          isScalarOperation(workload, op);
 }

 // Form dispatch regions from slice of the operation.
 static FailureOr<IREE::Flow::DispatchRegionOp>
 formDispatchRegionFromSlice(RewriterBase &rewriter, Operation *rootOp,
                             ArrayRef<Operation *> slice) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(rootOp);
   FailureOr<IREE::Flow::DispatchRegionOp> dispatchRegionOp =
       IREE::Flow::wrapOpInDispatchRegion(rewriter, rootOp);
   if (failed(dispatchRegionOp)) {
     return rootOp->emitOpError("failed to form dispatch region with root op");
   }
   FailureOr<IREE::Flow::DispatchRegionOp> newDispatchOp =
       movePrecedingOpsIntoDispatchRegion(rewriter, slice,
                                          dispatchRegionOp.value());
   if (failed(newDispatchOp)) {
     return dispatchRegionOp.value()->emitOpError(
         "failed to move slice into op");
   }
   return newDispatchOp.value();
 }

 void FormScalarDispatchesPass::runOnOperation() {
   mlir::FunctionOpInterface funcOp = getOperation();
   MLIRContext *context = &getContext();

   int scalarWorkloadLimit = 1;
   // Convenient struct to hold all operations that need to be moved into a
   // descriptor.
   struct DispatchRegionDescriptor {
     Operation *rootOp;
     SmallVector<Operation *> fusedOps;
   };

   SmallVector<DispatchRegionDescriptor> dispatches;
   llvm::DenseMap<Operation *, Operation *> opToRootMap;

   // Walk the function in postorder, reverse orded ignore all operations
   // not immediately nested within the `funcOp`.
   funcOp.walk<WalkOrder::PostOrder, ReverseIterator>([&](Operation *op) {
     if (op->getParentOp() != funcOp || opToRootMap.count(op)) {
       return;
     }

     if (!isSliceRoot(scalarWorkloadLimit, op)) {
       return;
     }

     llvm::SetVector<Operation *> fusedOpsSet =
         computeSliceToMoveIntoDispatch(scalarWorkloadLimit, op, opToRootMap);
     for (Operation *sliceOp : fusedOpsSet) {
       assert(!opToRootMap.count(sliceOp) &&
              "trying to add same op to two dispatches");
       opToRootMap[sliceOp] = op;
     }

     // Iterate backwards within the block to get ops that dont necessarily
     // have producer -> consumer relationship but can still be fused.
     Block *currBlock = op->getBlock();
     Operation *prevOp = op;
     bool didHorizontalFusion = false;
     llvm::SetVector<Operation *> ineligibleRoots;
     while (prevOp != &currBlock->front()) {
       prevOp = prevOp->getPrevNode();

       // If this operation is used by a operation we previously visited, but we
       // couldn't fuse it, stop.
       if (ineligibleRoots.contains(prevOp)) {
         break;
       }

       if (opToRootMap.count(prevOp)) {
         continue;
       }

       if (!isSliceRoot(scalarWorkloadLimit, prevOp)) {
         if (fusedOpsSet.contains(prevOp)) {
           continue;
         }
         // If this op is not being fused, any operations that defines values
         // used by this op cannot be horizontally fused
         // Insert all operations into the set that define op's operands or
         // define values used inside of op's regions
         mlir::visitUsedValuesDefinedAbove(
             prevOp->getRegions(), [&](OpOperand *operand) {
               if (auto definingOp = operand->get().getDefiningOp()) {
                 ineligibleRoots.insert(definingOp);
               }
             });

         for (Value val : prevOp->getOperands()) {
           if (auto definingOp = val.getDefiningOp()) {
             ineligibleRoots.insert(definingOp);
           }
         }
         continue;
       }

       didHorizontalFusion = true;
       fusedOpsSet.insert(prevOp);
       opToRootMap[prevOp] = op;
       llvm::SetVector<Operation *> currSlice = computeSliceToMoveIntoDispatch(
           scalarWorkloadLimit, prevOp, opToRootMap);
       for (auto sliceOp : currSlice) {
         assert(!opToRootMap.count(sliceOp) &&
                "trying to add same op to two dispatches");
         opToRootMap[sliceOp] = op;
       }
       fusedOpsSet.insert(currSlice.begin(), currSlice.end());
     }

     DispatchRegionDescriptor &currDispatch =
         dispatches.emplace_back(DispatchRegionDescriptor{});
     currDispatch.rootOp = op;
     currDispatch.fusedOps.assign(fusedOpsSet.begin(), fusedOpsSet.end());
     if (didHorizontalFusion) {
       mlir::computeTopologicalSorting(currDispatch.fusedOps);
     }
   });

   LLVM_DEBUG({
     llvm::dbgs() << "Num scalar dispatches : " << dispatches.size() << "\n";
     for (auto [index, dispatch] : llvm::enumerate(dispatches)) {
       llvm::dbgs() << "//--------------------------//\n";
       llvm::dbgs() << "Dispatch : " << index << ", Root :";
       dispatch.rootOp->print(llvm::dbgs());
       llvm::dbgs() << "\nFusedOps :";
       for (auto fusedOp : dispatch.fusedOps) {
         fusedOp->print(llvm::dbgs());
         llvm::dbgs() << "\n";
       }
       llvm::dbgs() << "//--------------------------//\n";
     }
   });

   IRRewriter rewriter(context);
   for (auto &currDispatch : dispatches) {
     rewriter.setInsertionPoint(currDispatch.rootOp);
     FailureOr<IREE::Flow::DispatchRegionOp> dispatchRegionOp =
         formDispatchRegionFromSlice(rewriter, currDispatch.rootOp,
                                     currDispatch.fusedOps);
     if (failed(dispatchRegionOp)) {
       currDispatch.rootOp->emitOpError(
           "failed to form scalar dispatch region with operation as root");
       return signalPassFailure();
     }

     // Set the workgroup count to {1, 1, 1} since this is to be executed
     // sequentially (at leats for now)
     Region &countRegion = dispatchRegionOp->getWorkgroupCount();
     Block *countBody = rewriter.createBlock(&countRegion, countRegion.begin());
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPointToStart(countBody);
     auto one = rewriter.create<arith::ConstantIndexOp>(
         dispatchRegionOp.value()->getLoc(), 1);
     rewriter.create<IREE::Flow::ReturnOp>(dispatchRegionOp.value()->getLoc(),
                                           ValueRange{one, one, one});
   }
 }

 } // namespace mlir::iree_compiler::DispatchCreation
	// Copyright 2023 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/Flow/IR/FlowOps.h"
	#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
	#include "iree/compiler/DispatchCreation/Passes.h"
	#include "llvm/ADT/SetVector.h"
	#include "llvm/Support/Debug.h"
	#include "mlir/Analysis/SliceAnalysis.h"
	#include "mlir/Analysis/TopologicalSortUtils.h"
	#include "mlir/Dialect/Affine/IR/AffineOps.h"
	#include "mlir/Dialect/Linalg/IR/Linalg.h"
	#include "mlir/Dialect/Tensor/IR/Tensor.h"
	#include "mlir/IR/Iterators.h"
	#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
	#include "mlir/Transforms/RegionUtils.h"

	#define DEBUG_TYPE "iree-dispatch-creation-form-scalar-dispatches"

	namespace mlir::iree_compiler::DispatchCreation {

	#define GEN_PASS_DEF_FORMSCALARDISPATCHESPASS
	#include "iree/compiler/DispatchCreation/Passes.h.inc"

	namespace {

	/// Pass declaration.
	struct FormScalarDispatchesPass final
	: public impl::FormScalarDispatchesPassBase<FormScalarDispatchesPass> {
	void runOnOperation() override;
	};
	} // namespace

	/// Return true if type represents a value less than `n` elements.
	static bool isScalarOrTensorOfLinearSizeN(int n, Type type) {
	if (type.isIntOrIndexOrFloat()) {
	return true;
	}
	if (auto tensorType = dyn_cast<RankedTensorType>(type)) {
	if (!tensorType.hasStaticShape()) {
	return false;
	}
	return tensorType.getNumElements() <= n;
	}
	return false;
	}

	/// Return `true` for operations that are to be treated as compute operations.
	static bool isComputeOperation(Operation *op) {
	MLIRContext *context = op->getContext();
	if (op->getDialect() == context->getLoadedDialect<linalg::LinalgDialect>()) {
	return true;
	}
	if (op->getDialect() == context->getLoadedDialect<tensor::TensorDialect>()) {
	return !isa<tensor::CastOp, tensor::CollapseShapeOp, tensor::EmptyOp,
	tensor::ExpandShapeOp, tensor::PackOp, tensor::UnPackOp>(op);
	}
	return false;
	}

	/// Return `true` if the workload of this operation is less than `n`.
	static bool isOperationWorkloadLessThanSizeN(int n, Operation *candidateOp) {
	return llvm::all_of(candidateOp->getOperands(),
	[&](Value v) {
	return isScalarOrTensorOfLinearSizeN(n, v.getType());
	}) &&
	llvm::all_of(candidateOp->getResultTypes(), [&](Type t) {
	return isScalarOrTensorOfLinearSizeN(n, t);
	});
	}

	/// Return `true` is the operation is to be treated as a scalar operation
	/// and moved into a scalar dispatch (not necessarily as the root of the
	/// dispatch).
	static bool isScalarOperation(int workload, Operation *op) {
	// 1. Ignore most operations. Only look for a whitelist set of operations.
	if (!isComputeOperation(op)) {
	return false;
	}

	// 2. Check that the workload of the operation is less then the limit
	if (!isOperationWorkloadLessThanSizeN(workload, op)) {
	return false;
	}

	// 3. Do not move operations that are cloned into the dispatch region.
	// TODO: This might prevent moving all scalar operations into dispatch
	// resulting in artifical splits. Revisit after more examples.
	return !IREE::Flow::isClonableIntoDispatchOp(op);
	}

	/// Given a `rootOp` return a DAG of the program that represents
	/// operations that can be moved into a scalar dispatch with the `rootOp`
	/// as the root of the DAG.
	llvm::SetVector<Operation *> computeSliceToMoveIntoDispatch(
	int workload, Operation *rootOp,
	const llvm::DenseMap<Operation , Operation > &opToRootMap) {
	BackwardSliceOptions options;
	options.filter = [&](Operation *currentOp) {
	assert(currentOp && "current op is null");
	if (opToRootMap.count(currentOp)) {
	return false;
	}
	// Operations needs to be in the same block as `rootOp`.
	if (currentOp->getBlock() != rootOp->getBlock()) {
	return false;
	}

	if (!isScalarOperation(workload, currentOp)) {
	return false;
	}

	// All its uses must be in the `opToRootMap`, i.e. they are either
	// in the current dispatches, or those already formed.
	return llvm::all_of(currentOp->getUsers(), [&](Operation *user) {
	return opToRootMap.count(user);
	});
	};
	options.omitBlockArguments = true;
	llvm::SetVector<Operation *> slice;
	getBackwardSlice(rootOp, &slice, options);
	return slice;
	}

	/// Return `true` if the op is to be treated as a root of a scalar dispatch.
	static bool isSliceRoot(int workload, Operation *op) {
	return !op->getParentOfType<IREE::Flow::DispatchRegionOp>() &&
	isScalarOperation(workload, op);
	}

	// Form dispatch regions from slice of the operation.
	static FailureOr<IREE::Flow::DispatchRegionOp>
	formDispatchRegionFromSlice(RewriterBase &rewriter, Operation *rootOp,
	ArrayRef<Operation *> slice) {
	OpBuilder::InsertionGuard g(rewriter);
	rewriter.setInsertionPoint(rootOp);
	FailureOr<IREE::Flow::DispatchRegionOp> dispatchRegionOp =
	IREE::Flow::wrapOpInDispatchRegion(rewriter, rootOp);
	if (failed(dispatchRegionOp)) {
	return rootOp->emitOpError("failed to form dispatch region with root op");
	}
	FailureOr<IREE::Flow::DispatchRegionOp> newDispatchOp =
	movePrecedingOpsIntoDispatchRegion(rewriter, slice,
	dispatchRegionOp.value());
	if (failed(newDispatchOp)) {
	return dispatchRegionOp.value()->emitOpError(
	"failed to move slice into op");
	}
	return newDispatchOp.value();
	}

	void FormScalarDispatchesPass::runOnOperation() {
	mlir::FunctionOpInterface funcOp = getOperation();
	MLIRContext *context = &getContext();

	int scalarWorkloadLimit = 1;
	// Convenient struct to hold all operations that need to be moved into a
	// descriptor.
	struct DispatchRegionDescriptor {
	Operation *rootOp;
	SmallVector<Operation *> fusedOps;
	};

	SmallVector<DispatchRegionDescriptor> dispatches;
	llvm::DenseMap<Operation , Operation > opToRootMap;

	// Walk the function in postorder, reverse orded ignore all operations
	// not immediately nested within the `funcOp`.
	funcOp.walk<WalkOrder::PostOrder, ReverseIterator>([&](Operation *op) {
	if (op->getParentOp() != funcOp \|\| opToRootMap.count(op)) {
	return;
	}

	if (!isSliceRoot(scalarWorkloadLimit, op)) {
	return;
	}

	llvm::SetVector<Operation *> fusedOpsSet =
	computeSliceToMoveIntoDispatch(scalarWorkloadLimit, op, opToRootMap);
	for (Operation *sliceOp : fusedOpsSet) {
	assert(!opToRootMap.count(sliceOp) &&
	"trying to add same op to two dispatches");
	opToRootMap[sliceOp] = op;
	}

	// Iterate backwards within the block to get ops that dont necessarily
	// have producer -> consumer relationship but can still be fused.
	Block *currBlock = op->getBlock();
	Operation *prevOp = op;
	bool didHorizontalFusion = false;
	llvm::SetVector<Operation *> ineligibleRoots;
	while (prevOp != &currBlock->front()) {
	prevOp = prevOp->getPrevNode();

	// If this operation is used by a operation we previously visited, but we
	// couldn't fuse it, stop.
	if (ineligibleRoots.contains(prevOp)) {
	break;
	}

	if (opToRootMap.count(prevOp)) {
	continue;
	}

	if (!isSliceRoot(scalarWorkloadLimit, prevOp)) {
	if (fusedOpsSet.contains(prevOp)) {
	continue;
	}
	// If this op is not being fused, any operations that defines values
	// used by this op cannot be horizontally fused
	// Insert all operations into the set that define op's operands or
	// define values used inside of op's regions
	mlir::visitUsedValuesDefinedAbove(
	prevOp->getRegions(), [&](OpOperand *operand) {
	if (auto definingOp = operand->get().getDefiningOp()) {
	ineligibleRoots.insert(definingOp);
	}
	});

	for (Value val : prevOp->getOperands()) {
	if (auto definingOp = val.getDefiningOp()) {
	ineligibleRoots.insert(definingOp);
	}
	}
	continue;
	}

	didHorizontalFusion = true;
	fusedOpsSet.insert(prevOp);
	opToRootMap[prevOp] = op;
	llvm::SetVector<Operation *> currSlice = computeSliceToMoveIntoDispatch(
	scalarWorkloadLimit, prevOp, opToRootMap);
	for (auto sliceOp : currSlice) {
	assert(!opToRootMap.count(sliceOp) &&
	"trying to add same op to two dispatches");
	opToRootMap[sliceOp] = op;
	}
	fusedOpsSet.insert(currSlice.begin(), currSlice.end());
	}

	DispatchRegionDescriptor &currDispatch =
	dispatches.emplace_back(DispatchRegionDescriptor{});
	currDispatch.rootOp = op;
	currDispatch.fusedOps.assign(fusedOpsSet.begin(), fusedOpsSet.end());
	if (didHorizontalFusion) {
	mlir::computeTopologicalSorting(currDispatch.fusedOps);
	}
	});

	LLVM_DEBUG({
	llvm::dbgs() << "Num scalar dispatches : " << dispatches.size() << "\n";
	for (auto [index, dispatch] : llvm::enumerate(dispatches)) {
	llvm::dbgs() << "//--------------------------//\n";
	llvm::dbgs() << "Dispatch : " << index << ", Root :";
	dispatch.rootOp->print(llvm::dbgs());
	llvm::dbgs() << "\nFusedOps :";
	for (auto fusedOp : dispatch.fusedOps) {
	fusedOp->print(llvm::dbgs());
	llvm::dbgs() << "\n";
	}
	llvm::dbgs() << "//--------------------------//\n";
	}
	});

	IRRewriter rewriter(context);
	for (auto &currDispatch : dispatches) {
	rewriter.setInsertionPoint(currDispatch.rootOp);
	FailureOr<IREE::Flow::DispatchRegionOp> dispatchRegionOp =
	formDispatchRegionFromSlice(rewriter, currDispatch.rootOp,
	currDispatch.fusedOps);
	if (failed(dispatchRegionOp)) {
	currDispatch.rootOp->emitOpError(
	"failed to form scalar dispatch region with operation as root");
	return signalPassFailure();
	}

	// Set the workgroup count to {1, 1, 1} since this is to be executed
	// sequentially (at leats for now)
	Region &countRegion = dispatchRegionOp->getWorkgroupCount();
	Block *countBody = rewriter.createBlock(&countRegion, countRegion.begin());
	OpBuilder::InsertionGuard g(rewriter);
	rewriter.setInsertionPointToStart(countBody);
	auto one = rewriter.create<arith::ConstantIndexOp>(
	dispatchRegionOp.value()->getLoc(), 1);
	rewriter.create<IREE::Flow::ReturnOp>(dispatchRegionOp.value()->getLoc(),
	ValueRange{one, one, one});
	}
	}

	} // namespace mlir::iree_compiler::DispatchCreation