compiler/src/iree/compiler/Modules/HAL/Inline/Transforms/InlineExecutables.cpp - 3p/openxla/iree - Git at Google

 // Copyright 2022 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/HAL/IR/HALDialect.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Modules/HAL/Inline/IR/HALInlineDialect.h"
 #include "iree/compiler/Modules/HAL/Inline/Transforms/Passes.h"
 #include "iree/compiler/Utils/IntegerSet.h"
 #include "iree/compiler/Utils/ModuleUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"

 namespace mlir::iree_compiler::IREE::HAL::Inline {

 #define GEN_PASS_DEF_INLINEEXECUTABLESPASS
 #include "iree/compiler/Modules/HAL/Inline/Transforms/Passes.h.inc"

 namespace {

 class InlineExecutablesPass final
     : public impl::InlineExecutablesPassBase<InlineExecutablesPass> {
 public:
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::Util::UtilDialect, IREE::HAL::HALDialect,
                     IREE::HAL::Inline::HALInlineDialect, arith::ArithDialect,
                     func::FuncDialect, scf::SCFDialect>();
   }

   void runOnOperation() override {
     mlir::ModuleOp moduleOp = getOperation();

     // Inline variants and produce a function map.
     DenseMap<Attribute, Attribute> exportToFuncMap;
     SymbolTableCollection symbolTables;
     for (auto executableOp : llvm::make_early_inc_range(
              moduleOp.getOps<IREE::HAL::ExecutableOp>())) {
       // Inline each variant.
       for (auto variantOp :
            executableOp.getOps<IREE::HAL::ExecutableVariantOp>()) {
         if (failed(inlineVariant(executableOp, variantOp, moduleOp,
                                  exportToFuncMap, symbolTables))) {
           return signalPassFailure();
         }
       }

       // Drop executable after information has been extracted and the workgroup
       // code has been inlined.
       executableOp.erase();
     }

     // Annotate all dispatches with the target function.
     for (auto funcOp : moduleOp.getOps<mlir::FunctionOpInterface>()) {
       auto result = funcOp.walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
         // Specify new target function that conversion can use to make the call.
         // We only support single variant dispatches when inline.
         auto entryPointAttrs = dispatchOp.getEntryPoints().getValue();
         if (entryPointAttrs.size() != 1) {
           dispatchOp.emitOpError()
               << "multiple variant targets not supported with the inline HAL";
           return WalkResult::interrupt();
         }
         auto targetFuncName =
             cast<StringAttr>(exportToFuncMap[entryPointAttrs.front()]);
         assert(targetFuncName && "missing mapping");
         dispatchOp->setAttr("hal_inline.target",
                             FlatSymbolRefAttr::get(targetFuncName));
         return WalkResult::advance();
       });
       if (result.wasInterrupted()) {
         return signalPassFailure();
       }
     }
   }

   LogicalResult inlineVariant(IREE::HAL::ExecutableOp executableOp,
                               IREE::HAL::ExecutableVariantOp variantOp,
                               mlir::ModuleOp targetModuleOp,
                               DenseMap<Attribute, Attribute> &exportToFuncMap,
                               SymbolTableCollection &symbolTables) {
     auto innerModuleOp = variantOp.getInnerModule();
     auto innerSymbolTable = symbolTables.getSymbolTable(innerModuleOp);
     auto innerModuleBuilder = OpBuilder::atBlockEnd(innerModuleOp.getBody());

     // We want to merge the module ahead of the exported functions to ensure
     // initializer order is preserved.
     OpBuilder targetModuleBuilder(executableOp);

     // Build each dispatch function wrapper.
     auto indexType = innerModuleBuilder.getIndexType();
     auto i32Type = innerModuleBuilder.getI32Type();
     auto bufferType = innerModuleBuilder.getType<IREE::Util::BufferType>();
     for (auto exportOp : variantOp.getExportOps()) {
       // Build dispatch function signature that the stream.cmd.dispatch ops will
       // map to.
       auto layoutAttr = exportOp.getLayout();
       size_t bindingCount = layoutAttr.getBindings().size();
       SmallVector<Type> inputTypes;
       inputTypes.append(exportOp.getWorkgroupCountBody()->getNumArguments() - 1,
                         indexType); // workload
       inputTypes.append(layoutAttr.getConstants(), i32Type);
       inputTypes.append(bindingCount, bufferType); // buffers
       inputTypes.append(bindingCount, indexType);  // offsets
       inputTypes.append(bindingCount, indexType);  // lengths
       auto dispatchFuncType =
           innerModuleBuilder.getFunctionType(inputTypes, {});

       // Create the function and insert into the module.
       auto dispatchFuncOp = IREE::Util::FuncOp::create(
           exportOp.getLoc(),
           ("__dispatch_" + executableOp.getName() + "_" + exportOp.getName())
               .str(),
           dispatchFuncType);
       dispatchFuncOp.setPrivate();
       innerSymbolTable.insert(dispatchFuncOp,
                               innerModuleBuilder.getInsertionPoint());
       innerModuleBuilder.setInsertionPointAfter(dispatchFuncOp);

       // Build the dispatch function by calling the target function in a loop.
       auto bodyFuncOp =
           innerSymbolTable.lookup<FunctionOpInterface>(exportOp.getName());
       if (!bodyFuncOp) {
         return exportOp.emitOpError("missing body function");
       }
       if (bodyFuncOp.isPublic()) {
         if (failed(rewriteWorkgroupSignature(layoutAttr, bindingCount,
                                              bodyFuncOp))) {
           return failure();
         }
         bodyFuncOp.setPrivate(); // so we only do it once
       }
       buildDispatchFunc(exportOp, layoutAttr, bindingCount, bodyFuncOp,
                         dispatchFuncOp);

       // Map from what the stream.cmd.dispatch ops is using to the new function.
       auto exportTargetAttr = SymbolRefAttr::get(
           executableOp.getNameAttr(),
           {
               FlatSymbolRefAttr::get(variantOp.getNameAttr()),
               FlatSymbolRefAttr::get(exportOp.getNameAttr()),
           });
       exportToFuncMap[exportTargetAttr] = dispatchFuncOp.getNameAttr();
     }

     // Merge the source executable module into the target host module.
     if (failed(mergeModuleInto(innerModuleOp, targetModuleOp,
                                targetModuleBuilder))) {
       return failure();
     }

     return success();
   }

   // Rewrites a workgroup body function signature to a flattened list.
   //
   // Body (as translated):
   //   (local_memory, [constants], [bindings],
   //    workgroup_x, workgroup_y, workgroup_z,
   //    workgroup_size_x, workgroup_size_y, workgroup_size_z,
   //    workgroup_count_x, workgroup_count_y, workgroup_count_z)
   //
   // Body after rewrite:
   //   (local_memory, constants..., bindings...,
   //    workgroup_x, workgroup_y, workgroup_z,
   //    workgroup_size_x, workgroup_size_y, workgroup_size_z,
   //    workgroup_count_x, workgroup_count_y, workgroup_count_z)
   //
   // To make this process easier and lighten the load on the downstream passes
   // we muck with the ABI to pass a flattened list of constants and bindings.
   // Whenever better IPO and util.list optimizations are added we could back
   // this out to keep things vanilla and have fewer places making assumptions
   // about the function signatures.
   LogicalResult
   rewriteWorkgroupSignature(IREE::HAL::PipelineLayoutAttr layoutAttr,
                             size_t bindingCount,
                             FunctionOpInterface bodyFuncOp) {
     auto *entryBlock = &bodyFuncOp.front();
     auto builder = OpBuilder::atBlockBegin(entryBlock);
     auto indexType = builder.getIndexType();
     auto i32Type = builder.getI32Type();
     auto bufferType = builder.getType<IREE::Util::BufferType>();

     // There may be nicer ways of doing this but I can't find them.
     // We build a new list of argument types and insert them as we go. This lets
     // us map the arguments over and replace usage such that by the end we can
     // slice off the original arguments as they'll have no more uses.
     unsigned originalArgCount = entryBlock->getNumArguments();
     SmallVector<Type> newArgTypes;
     unsigned argOffset = 0;

     // Local memory is carried across as-is.
     auto localMemoryArg = entryBlock->getArgument(argOffset++);
     newArgTypes.push_back(bufferType);
     localMemoryArg.replaceAllUsesWith(
         entryBlock->addArgument(bufferType, localMemoryArg.getLoc()));

     // Expand push constants by replacing buffer accesses with the flattened
     // args.
     newArgTypes.append(layoutAttr.getConstants(), i32Type);
     auto constantBuffer = entryBlock->getArgument(argOffset++);
     SmallVector<Value> constantArgs;
     for (unsigned i = 0; i < layoutAttr.getConstants(); ++i) {
       constantArgs.push_back(
           entryBlock->addArgument(i32Type, constantBuffer.getLoc()));
     }
     if (failed(replaceBufferAccesses(constantBuffer, constantArgs))) {
       return failure();
     }

     // Expand buffer list by replacing list accesses with the flattened args.
     newArgTypes.append(bindingCount, bufferType);
     auto bindingList = entryBlock->getArgument(argOffset++);
     SmallVector<Value> bindingArgs;
     for (unsigned i = 0; i < bindingCount; ++i) {
       bindingArgs.push_back(
           entryBlock->addArgument(bufferType, bindingList.getLoc()));
     }
     if (failed(replaceListAccesses(bindingList, bindingArgs))) {
       return failure();
     }

     // Take care of the workgroup id/size/count tuples.
     auto entryBuilder = OpBuilder::atBlockBegin(entryBlock);
     for (unsigned i = 0; i < 3 * /*xyz=*/3; ++i) {
       newArgTypes.push_back(indexType);
       auto oldArg = entryBlock->getArgument(argOffset++);
       auto newArg = entryBlock->addArgument(indexType, oldArg.getLoc());
       oldArg.replaceAllUsesWith(arith::IndexCastOp::create(
           entryBuilder, oldArg.getLoc(), i32Type, newArg));
     }

     // Erase the original args.
     for (unsigned i = 0; i < originalArgCount; ++i) {
       entryBlock->eraseArgument(0);
     }

     // Update function signature to reflect the entry block args.
     bodyFuncOp.setType(
         builder.getFunctionType(newArgTypes, bodyFuncOp.getResultTypes()));

     return success();
   }

   // Replaces trivial constant index accesses to a buffer with their values.
   // This is an extremely poor optimization that we should remove if buffer
   // ever gets store-load forwarding - we could just create the buffer, store
   // the elements, and let that take care of the rest. Today it doesn't do that.
   LogicalResult replaceBufferAccesses(Value buffer, ValueRange elements) {
     for (auto user : llvm::make_early_inc_range(buffer.getUsers())) {
       if (auto sizeOp = dyn_cast<IREE::Util::BufferSizeOp>(user)) {
         // Ignored but we need to get rid of it.
         // TODO(benvanik): see if we can allow this through; today it will pin
         // the function argument (constants most likely) and cause us to fail to
         // remove it later on.
         OpBuilder builder(sizeOp);
         Value dummySize = arith::ConstantIndexOp::create(
             builder, sizeOp.getLoc(), 0xCAFEF00D);
         sizeOp.replaceAllUsesWith(dummySize);
         sizeOp.erase();
         continue;
       } else if (auto loadOp = dyn_cast<IREE::Util::BufferLoadOp>(user)) {
         APInt index;
         if (matchPattern(loadOp.getSourceOffset(), m_ConstantInt(&index))) {
           loadOp.replaceAllUsesWith(
               elements[index.getSExtValue() / sizeof(uint32_t)]);
           loadOp.erase();
           continue;
         } else {
           return loadOp.emitOpError(
               "unhandled dynamic buffer access; must be static");
         }
       } else if (auto loadOp = dyn_cast<memref::LoadOp>(user)) {
         if (loadOp.getIndices().size() != 1) {
           return loadOp.emitOpError(
               "expected memrefs to have been flattened before inlining "
               "executables");
         }
         APInt index;
         if (matchPattern(loadOp.getIndices()[0], m_ConstantInt(&index))) {
           loadOp.replaceAllUsesWith(elements[index.getSExtValue()]);
           loadOp.erase();
           continue;
         } else {
           return loadOp.emitOpError(
               "unhandled dynamic buffer access; must be static");
         }
       } else {
         return user->emitOpError(
             "unhandled buffer access op; only loads are supported");
       }
     }
     return success();
   }

   // Replaces trivial constant index accesses to a list with their values.
   // util.list store-load forwarding could do this instead.
   LogicalResult replaceListAccesses(Value list, ValueRange elements) {
     for (auto user : llvm::make_early_inc_range(list.getUsers())) {
       if (auto getOp = dyn_cast<IREE::Util::ListGetOp>(user)) {
         APInt index;
         if (matchPattern(getOp.getIndex(), m_ConstantInt(&index))) {
           getOp.replaceAllUsesWith(elements[index.getSExtValue()]);
           getOp.erase();
           continue;
         } else {
           return getOp.emitOpError(
               "unhandled dynamic list access; must be static");
         }
       } else {
         return user->emitOpError(
             "unhandled list access op; only gets are supported");
       }
     }
     return success();
   }

   // Builds a function that calls a workgroup body and marshals arguments.
   //
   // Incoming:
   //   (workload..., constants...,
   //    binding_buffers..., binding_offsets..., binding_lengths...)
   // Body (as translated):
   //   (local_memory, [constants], [bindings],
   //    workgroup_x, workgroup_y, workgroup_z,
   //    workgroup_size_x, workgroup_size_y, workgroup_size_z,
   //    workgroup_count_x, workgroup_count_y, workgroup_count_z)
   void buildDispatchFunc(IREE::HAL::ExecutableExportOp exportOp,
                          IREE::HAL::PipelineLayoutAttr layoutAttr,
                          size_t bindingCount, FunctionOpInterface bodyFuncOp,
                          FunctionOpInterface dispatchFuncOp) {
     auto loc = exportOp.getLoc();
     auto builder = OpBuilder::atBlockBegin(dispatchFuncOp.addEntryBlock());
     IndexSet indexSet(loc, builder);
     auto bufferType = builder.getType<IREE::Util::BufferType>();

     SmallVector<Value> workgroupArgs;

     // Calculate the XYZ workgroup count from the export function.
     // There may be multiple exports pointing at the same body with different
     // workgroup count functions.
     unsigned workloadArgCount =
         exportOp.getWorkgroupCountBody()->getNumArguments() - 1;
     unsigned argOffset = 0;
     SmallVector<Value> workload;
     workload.reserve(workloadArgCount);
     for (unsigned i = 0; i < workloadArgCount; ++i) {
       workload.push_back(dispatchFuncOp.getArgument(argOffset++));
     }
     Value device = IREE::Util::NullOp::create(
         builder, loc, builder.getType<IREE::HAL::DeviceType>());
     std::array<Value, 3> workgroupCount =
         exportOp.calculateWorkgroupCount(loc, device, workload, builder);

     // For now we don't handle local memory.
     Value localMemory = IREE::Util::NullOp::create(builder, loc, bufferType);
     workgroupArgs.push_back(localMemory);

     // Pass all constants through.
     for (int64_t i = 0; i < layoutAttr.getConstants(); ++i) {
       workgroupArgs.push_back(dispatchFuncOp.getArgument(argOffset++));
     }

     // Pass all buffers through as subspans with the binding offset and length
     // factored in. IPO can propagate the subspans (hopefully).
     for (size_t i = 0; i < bindingCount; ++i) {
       auto bindingBuffer = dispatchFuncOp.getArgument(argOffset + i);
       auto bindingOffset =
           dispatchFuncOp.getArgument(argOffset + bindingCount + i);
       auto bindingLength = dispatchFuncOp.getArgument(argOffset + bindingCount +
                                                       bindingCount + i);
       Value bufferSize =
           IREE::Util::BufferSizeOp::create(builder, loc, bindingBuffer);
       Value bindingView = IREE::Util::BufferSubspanOp::create(
           builder, loc, bindingBuffer, bufferSize, bindingOffset,
           bindingLength);
       workgroupArgs.push_back(bindingView);
     }

     int workgroupXYZOffset = workgroupArgs.size();
     workgroupArgs.append(3, nullptr);         // workgroup_xyz, set below
     workgroupArgs.append(3, indexSet.get(1)); // workgroup_size_xyz
     llvm::append_range(workgroupArgs, workgroupCount); // workgroup_count_xyz

     // Z -> Y -> Z loop nest.
     scf::ForOp::create(
         builder, loc, indexSet.get(0), workgroupCount[2], indexSet.get(1),
         ValueRange{},
         [&](OpBuilder &forZBuilder, Location loc, Value iz, ValueRange iters) {
           workgroupArgs[workgroupXYZOffset + 2] = iz;
           scf::ForOp::create(
               forZBuilder, loc, indexSet.get(0), workgroupCount[1],
               indexSet.get(1), ValueRange{},
               [&](OpBuilder &forYBuilder, Location loc, Value iy,
                   ValueRange iters) {
                 workgroupArgs[workgroupXYZOffset + 1] = iy;
                 scf::ForOp::create(
                     forYBuilder, loc, indexSet.get(0), workgroupCount[0],
                     indexSet.get(1), ValueRange{},
                     [&](OpBuilder &forXBuilder, Location loc, Value ix,
                         ValueRange iters) {
                       workgroupArgs[workgroupXYZOffset + 0] = ix;
                       func::CallOp::create(
                           forXBuilder, loc, bodyFuncOp.getNameAttr(),
                           bodyFuncOp.getResultTypes(), workgroupArgs);
                       scf::YieldOp::create(forXBuilder, loc);
                     });
                 scf::YieldOp::create(forYBuilder, loc);
               });
           scf::YieldOp::create(forZBuilder, loc);
         });

     IREE::Util::ReturnOp::create(builder, loc);
   }
 };

 } // namespace
 } // namespace mlir::iree_compiler::IREE::HAL::Inline
	// Copyright 2022 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/HAL/IR/HALDialect.h"
	#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
	#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
	#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
	#include "iree/compiler/Modules/HAL/Inline/IR/HALInlineDialect.h"
	#include "iree/compiler/Modules/HAL/Inline/Transforms/Passes.h"
	#include "iree/compiler/Utils/IntegerSet.h"
	#include "iree/compiler/Utils/ModuleUtils.h"
	#include "llvm/ADT/STLExtras.h"
	#include "mlir/Dialect/Arith/IR/Arith.h"
	#include "mlir/Dialect/Func/IR/FuncOps.h"
	#include "mlir/Dialect/MemRef/IR/MemRef.h"
	#include "mlir/Dialect/SCF/IR/SCF.h"
	#include "mlir/IR/BuiltinOps.h"
	#include "mlir/Pass/Pass.h"
	#include "mlir/Pass/PassRegistry.h"

	namespace mlir::iree_compiler::IREE::HAL::Inline {

	#define GEN_PASS_DEF_INLINEEXECUTABLESPASS
	#include "iree/compiler/Modules/HAL/Inline/Transforms/Passes.h.inc"

	namespace {

	class InlineExecutablesPass final
	: public impl::InlineExecutablesPassBase<InlineExecutablesPass> {
	public:
	void getDependentDialects(DialectRegistry &registry) const override {
	registry.insert<IREE::Util::UtilDialect, IREE::HAL::HALDialect,
	IREE::HAL::Inline::HALInlineDialect, arith::ArithDialect,
	func::FuncDialect, scf::SCFDialect>();
	}

	void runOnOperation() override {
	mlir::ModuleOp moduleOp = getOperation();

	// Inline variants and produce a function map.
	DenseMap<Attribute, Attribute> exportToFuncMap;
	SymbolTableCollection symbolTables;
	for (auto executableOp : llvm::make_early_inc_range(
	moduleOp.getOps<IREE::HAL::ExecutableOp>())) {
	// Inline each variant.
	for (auto variantOp :
	executableOp.getOps<IREE::HAL::ExecutableVariantOp>()) {
	if (failed(inlineVariant(executableOp, variantOp, moduleOp,
	exportToFuncMap, symbolTables))) {
	return signalPassFailure();
	}
	}

	// Drop executable after information has been extracted and the workgroup
	// code has been inlined.
	executableOp.erase();
	}

	// Annotate all dispatches with the target function.
	for (auto funcOp : moduleOp.getOps<mlir::FunctionOpInterface>()) {
	auto result = funcOp.walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
	// Specify new target function that conversion can use to make the call.
	// We only support single variant dispatches when inline.
	auto entryPointAttrs = dispatchOp.getEntryPoints().getValue();
	if (entryPointAttrs.size() != 1) {
	dispatchOp.emitOpError()
	<< "multiple variant targets not supported with the inline HAL";
	return WalkResult::interrupt();
	}
	auto targetFuncName =
	cast<StringAttr>(exportToFuncMap[entryPointAttrs.front()]);
	assert(targetFuncName && "missing mapping");
	dispatchOp->setAttr("hal_inline.target",
	FlatSymbolRefAttr::get(targetFuncName));
	return WalkResult::advance();
	});
	if (result.wasInterrupted()) {
	return signalPassFailure();
	}
	}
	}

	LogicalResult inlineVariant(IREE::HAL::ExecutableOp executableOp,
	IREE::HAL::ExecutableVariantOp variantOp,
	mlir::ModuleOp targetModuleOp,
	DenseMap<Attribute, Attribute> &exportToFuncMap,
	SymbolTableCollection &symbolTables) {
	auto innerModuleOp = variantOp.getInnerModule();
	auto innerSymbolTable = symbolTables.getSymbolTable(innerModuleOp);
	auto innerModuleBuilder = OpBuilder::atBlockEnd(innerModuleOp.getBody());

	// We want to merge the module ahead of the exported functions to ensure
	// initializer order is preserved.
	OpBuilder targetModuleBuilder(executableOp);

	// Build each dispatch function wrapper.
	auto indexType = innerModuleBuilder.getIndexType();
	auto i32Type = innerModuleBuilder.getI32Type();
	auto bufferType = innerModuleBuilder.getType<IREE::Util::BufferType>();
	for (auto exportOp : variantOp.getExportOps()) {
	// Build dispatch function signature that the stream.cmd.dispatch ops will
	// map to.
	auto layoutAttr = exportOp.getLayout();
	size_t bindingCount = layoutAttr.getBindings().size();
	SmallVector<Type> inputTypes;
	inputTypes.append(exportOp.getWorkgroupCountBody()->getNumArguments() - 1,
	indexType); // workload
	inputTypes.append(layoutAttr.getConstants(), i32Type);
	inputTypes.append(bindingCount, bufferType); // buffers
	inputTypes.append(bindingCount, indexType); // offsets
	inputTypes.append(bindingCount, indexType); // lengths
	auto dispatchFuncType =
	innerModuleBuilder.getFunctionType(inputTypes, {});

	// Create the function and insert into the module.
	auto dispatchFuncOp = IREE::Util::FuncOp::create(
	exportOp.getLoc(),
	("__dispatch_" + executableOp.getName() + "_" + exportOp.getName())
	.str(),
	dispatchFuncType);
	dispatchFuncOp.setPrivate();
	innerSymbolTable.insert(dispatchFuncOp,
	innerModuleBuilder.getInsertionPoint());
	innerModuleBuilder.setInsertionPointAfter(dispatchFuncOp);

	// Build the dispatch function by calling the target function in a loop.
	auto bodyFuncOp =
	innerSymbolTable.lookup<FunctionOpInterface>(exportOp.getName());
	if (!bodyFuncOp) {
	return exportOp.emitOpError("missing body function");
	}
	if (bodyFuncOp.isPublic()) {
	if (failed(rewriteWorkgroupSignature(layoutAttr, bindingCount,
	bodyFuncOp))) {
	return failure();
	}
	bodyFuncOp.setPrivate(); // so we only do it once
	}
	buildDispatchFunc(exportOp, layoutAttr, bindingCount, bodyFuncOp,
	dispatchFuncOp);

	// Map from what the stream.cmd.dispatch ops is using to the new function.
	auto exportTargetAttr = SymbolRefAttr::get(
	executableOp.getNameAttr(),
	{
	FlatSymbolRefAttr::get(variantOp.getNameAttr()),
	FlatSymbolRefAttr::get(exportOp.getNameAttr()),
	});
	exportToFuncMap[exportTargetAttr] = dispatchFuncOp.getNameAttr();
	}

	// Merge the source executable module into the target host module.
	if (failed(mergeModuleInto(innerModuleOp, targetModuleOp,
	targetModuleBuilder))) {
	return failure();
	}

	return success();
	}

	// Rewrites a workgroup body function signature to a flattened list.
	//
	// Body (as translated):
	// (local_memory, [constants], [bindings],
	// workgroup_x, workgroup_y, workgroup_z,
	// workgroup_size_x, workgroup_size_y, workgroup_size_z,
	// workgroup_count_x, workgroup_count_y, workgroup_count_z)
	//
	// Body after rewrite:
	// (local_memory, constants..., bindings...,
	// workgroup_x, workgroup_y, workgroup_z,
	// workgroup_size_x, workgroup_size_y, workgroup_size_z,
	// workgroup_count_x, workgroup_count_y, workgroup_count_z)
	//
	// To make this process easier and lighten the load on the downstream passes
	// we muck with the ABI to pass a flattened list of constants and bindings.
	// Whenever better IPO and util.list optimizations are added we could back
	// this out to keep things vanilla and have fewer places making assumptions
	// about the function signatures.
	LogicalResult
	rewriteWorkgroupSignature(IREE::HAL::PipelineLayoutAttr layoutAttr,
	size_t bindingCount,
	FunctionOpInterface bodyFuncOp) {
	auto *entryBlock = &bodyFuncOp.front();
	auto builder = OpBuilder::atBlockBegin(entryBlock);
	auto indexType = builder.getIndexType();
	auto i32Type = builder.getI32Type();
	auto bufferType = builder.getType<IREE::Util::BufferType>();

	// There may be nicer ways of doing this but I can't find them.
	// We build a new list of argument types and insert them as we go. This lets
	// us map the arguments over and replace usage such that by the end we can
	// slice off the original arguments as they'll have no more uses.
	unsigned originalArgCount = entryBlock->getNumArguments();
	SmallVector<Type> newArgTypes;
	unsigned argOffset = 0;

	// Local memory is carried across as-is.
	auto localMemoryArg = entryBlock->getArgument(argOffset++);
	newArgTypes.push_back(bufferType);
	localMemoryArg.replaceAllUsesWith(
	entryBlock->addArgument(bufferType, localMemoryArg.getLoc()));

	// Expand push constants by replacing buffer accesses with the flattened
	// args.
	newArgTypes.append(layoutAttr.getConstants(), i32Type);
	auto constantBuffer = entryBlock->getArgument(argOffset++);
	SmallVector<Value> constantArgs;
	for (unsigned i = 0; i < layoutAttr.getConstants(); ++i) {
	constantArgs.push_back(
	entryBlock->addArgument(i32Type, constantBuffer.getLoc()));
	}
	if (failed(replaceBufferAccesses(constantBuffer, constantArgs))) {
	return failure();
	}

	// Expand buffer list by replacing list accesses with the flattened args.
	newArgTypes.append(bindingCount, bufferType);
	auto bindingList = entryBlock->getArgument(argOffset++);
	SmallVector<Value> bindingArgs;
	for (unsigned i = 0; i < bindingCount; ++i) {
	bindingArgs.push_back(
	entryBlock->addArgument(bufferType, bindingList.getLoc()));
	}
	if (failed(replaceListAccesses(bindingList, bindingArgs))) {
	return failure();
	}

	// Take care of the workgroup id/size/count tuples.
	auto entryBuilder = OpBuilder::atBlockBegin(entryBlock);
	for (unsigned i = 0; i < 3 * /xyz=/3; ++i) {
	newArgTypes.push_back(indexType);
	auto oldArg = entryBlock->getArgument(argOffset++);
	auto newArg = entryBlock->addArgument(indexType, oldArg.getLoc());
	oldArg.replaceAllUsesWith(arith::IndexCastOp::create(
	entryBuilder, oldArg.getLoc(), i32Type, newArg));
	}

	// Erase the original args.
	for (unsigned i = 0; i < originalArgCount; ++i) {
	entryBlock->eraseArgument(0);
	}

	// Update function signature to reflect the entry block args.
	bodyFuncOp.setType(
	builder.getFunctionType(newArgTypes, bodyFuncOp.getResultTypes()));

	return success();
	}

	// Replaces trivial constant index accesses to a buffer with their values.
	// This is an extremely poor optimization that we should remove if buffer
	// ever gets store-load forwarding - we could just create the buffer, store
	// the elements, and let that take care of the rest. Today it doesn't do that.
	LogicalResult replaceBufferAccesses(Value buffer, ValueRange elements) {
	for (auto user : llvm::make_early_inc_range(buffer.getUsers())) {
	if (auto sizeOp = dyn_cast<IREE::Util::BufferSizeOp>(user)) {
	// Ignored but we need to get rid of it.
	// TODO(benvanik): see if we can allow this through; today it will pin
	// the function argument (constants most likely) and cause us to fail to
	// remove it later on.
	OpBuilder builder(sizeOp);
	Value dummySize = arith::ConstantIndexOp::create(
	builder, sizeOp.getLoc(), 0xCAFEF00D);
	sizeOp.replaceAllUsesWith(dummySize);
	sizeOp.erase();
	continue;
	} else if (auto loadOp = dyn_cast<IREE::Util::BufferLoadOp>(user)) {
	APInt index;
	if (matchPattern(loadOp.getSourceOffset(), m_ConstantInt(&index))) {
	loadOp.replaceAllUsesWith(
	elements[index.getSExtValue() / sizeof(uint32_t)]);
	loadOp.erase();
	continue;
	} else {
	return loadOp.emitOpError(
	"unhandled dynamic buffer access; must be static");
	}
	} else if (auto loadOp = dyn_cast<memref::LoadOp>(user)) {
	if (loadOp.getIndices().size() != 1) {
	return loadOp.emitOpError(
	"expected memrefs to have been flattened before inlining "
	"executables");
	}
	APInt index;
	if (matchPattern(loadOp.getIndices()[0], m_ConstantInt(&index))) {
	loadOp.replaceAllUsesWith(elements[index.getSExtValue()]);
	loadOp.erase();
	continue;
	} else {
	return loadOp.emitOpError(
	"unhandled dynamic buffer access; must be static");
	}
	} else {
	return user->emitOpError(
	"unhandled buffer access op; only loads are supported");
	}
	}
	return success();
	}

	// Replaces trivial constant index accesses to a list with their values.
	// util.list store-load forwarding could do this instead.
	LogicalResult replaceListAccesses(Value list, ValueRange elements) {
	for (auto user : llvm::make_early_inc_range(list.getUsers())) {
	if (auto getOp = dyn_cast<IREE::Util::ListGetOp>(user)) {
	APInt index;
	if (matchPattern(getOp.getIndex(), m_ConstantInt(&index))) {
	getOp.replaceAllUsesWith(elements[index.getSExtValue()]);
	getOp.erase();
	continue;
	} else {
	return getOp.emitOpError(
	"unhandled dynamic list access; must be static");
	}
	} else {
	return user->emitOpError(
	"unhandled list access op; only gets are supported");
	}
	}
	return success();
	}

	// Builds a function that calls a workgroup body and marshals arguments.
	//
	// Incoming:
	// (workload..., constants...,
	// binding_buffers..., binding_offsets..., binding_lengths...)
	// Body (as translated):
	// (local_memory, [constants], [bindings],
	// workgroup_x, workgroup_y, workgroup_z,
	// workgroup_size_x, workgroup_size_y, workgroup_size_z,
	// workgroup_count_x, workgroup_count_y, workgroup_count_z)
	void buildDispatchFunc(IREE::HAL::ExecutableExportOp exportOp,
	IREE::HAL::PipelineLayoutAttr layoutAttr,
	size_t bindingCount, FunctionOpInterface bodyFuncOp,
	FunctionOpInterface dispatchFuncOp) {
	auto loc = exportOp.getLoc();
	auto builder = OpBuilder::atBlockBegin(dispatchFuncOp.addEntryBlock());
	IndexSet indexSet(loc, builder);
	auto bufferType = builder.getType<IREE::Util::BufferType>();

	SmallVector<Value> workgroupArgs;

	// Calculate the XYZ workgroup count from the export function.
	// There may be multiple exports pointing at the same body with different
	// workgroup count functions.
	unsigned workloadArgCount =
	exportOp.getWorkgroupCountBody()->getNumArguments() - 1;
	unsigned argOffset = 0;
	SmallVector<Value> workload;
	workload.reserve(workloadArgCount);
	for (unsigned i = 0; i < workloadArgCount; ++i) {
	workload.push_back(dispatchFuncOp.getArgument(argOffset++));
	}
	Value device = IREE::Util::NullOp::create(
	builder, loc, builder.getType<IREE::HAL::DeviceType>());
	std::array<Value, 3> workgroupCount =
	exportOp.calculateWorkgroupCount(loc, device, workload, builder);

	// For now we don't handle local memory.
	Value localMemory = IREE::Util::NullOp::create(builder, loc, bufferType);
	workgroupArgs.push_back(localMemory);

	// Pass all constants through.
	for (int64_t i = 0; i < layoutAttr.getConstants(); ++i) {
	workgroupArgs.push_back(dispatchFuncOp.getArgument(argOffset++));
	}

	// Pass all buffers through as subspans with the binding offset and length
	// factored in. IPO can propagate the subspans (hopefully).
	for (size_t i = 0; i < bindingCount; ++i) {
	auto bindingBuffer = dispatchFuncOp.getArgument(argOffset + i);
	auto bindingOffset =
	dispatchFuncOp.getArgument(argOffset + bindingCount + i);
	auto bindingLength = dispatchFuncOp.getArgument(argOffset + bindingCount +
	bindingCount + i);
	Value bufferSize =
	IREE::Util::BufferSizeOp::create(builder, loc, bindingBuffer);
	Value bindingView = IREE::Util::BufferSubspanOp::create(
	builder, loc, bindingBuffer, bufferSize, bindingOffset,
	bindingLength);
	workgroupArgs.push_back(bindingView);
	}

	int workgroupXYZOffset = workgroupArgs.size();
	workgroupArgs.append(3, nullptr); // workgroup_xyz, set below
	workgroupArgs.append(3, indexSet.get(1)); // workgroup_size_xyz
	llvm::append_range(workgroupArgs, workgroupCount); // workgroup_count_xyz

	// Z -> Y -> Z loop nest.
	scf::ForOp::create(
	builder, loc, indexSet.get(0), workgroupCount[2], indexSet.get(1),
	ValueRange{},
	[&](OpBuilder &forZBuilder, Location loc, Value iz, ValueRange iters) {
	workgroupArgs[workgroupXYZOffset + 2] = iz;
	scf::ForOp::create(
	forZBuilder, loc, indexSet.get(0), workgroupCount[1],
	indexSet.get(1), ValueRange{},
	[&](OpBuilder &forYBuilder, Location loc, Value iy,
	ValueRange iters) {
	workgroupArgs[workgroupXYZOffset + 1] = iy;
	scf::ForOp::create(
	forYBuilder, loc, indexSet.get(0), workgroupCount[0],
	indexSet.get(1), ValueRange{},
	[&](OpBuilder &forXBuilder, Location loc, Value ix,
	ValueRange iters) {
	workgroupArgs[workgroupXYZOffset + 0] = ix;
	func::CallOp::create(
	forXBuilder, loc, bodyFuncOp.getNameAttr(),
	bodyFuncOp.getResultTypes(), workgroupArgs);
	scf::YieldOp::create(forXBuilder, loc);
	});
	scf::YieldOp::create(forYBuilder, loc);
	});
	scf::YieldOp::create(forZBuilder, loc);
	});

	IREE::Util::ReturnOp::create(builder, loc);
	}
	};

	} // namespace
	} // namespace mlir::iree_compiler::IREE::HAL::Inline