[HAL] Add a pass to materialize homogeneous encodings.

compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodingPass.cpp

@@ -131,20 +131,33 @@ static MatmulTileParams chooseMatmulTileParams(EncodingUser user,
 
 struct CPUMaterializeEncodingPass
     : public CPUMaterializeEncodingBase<CPUMaterializeEncodingPass> {
+  CPUMaterializeEncodingPass() : targetAttr(nullptr) {}
+  explicit CPUMaterializeEncodingPass(IREE::HAL::ExecutableTargetAttr attr)
+      : targetAttr(attr) {}
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithDialect, IREE::LinalgExt::IREELinalgExtDialect,
                     IREE::Codegen::IREECodegenDialect>();
   }
   void runOnOperation() override;
+
+private:
+  IREE::HAL::ExecutableTargetAttr targetAttr;
 };
 
 struct CPUMaterializeUpperBoundTileSizePass
     : public CPUMaterializeUpperBoundTileSizeBase<
           CPUMaterializeUpperBoundTileSizePass> {
+  CPUMaterializeUpperBoundTileSizePass() = default;
+  explicit CPUMaterializeUpperBoundTileSizePass(
+      ArrayRef<IREE::HAL::ExecutableTargetAttr> attrs)
+      : targetAttrs(attrs) {}
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithDialect>();
   }
   void runOnOperation() override;
+
+private:
+  SmallVector<IREE::HAL::ExecutableTargetAttr, 4> targetAttrs;
 };
 
 FailureOr<MaterializeEncodingInfo>
@@ -235,7 +248,8 @@ void CPUMaterializeEncodingPass::runOnOperation() {
   MLIRContext *context = &getContext();
   auto operation = getOperation();
   RewritePatternSet materializeEncodingPattern(context);
-  auto targetAttr = ExecutableTargetAttr::lookup(operation);
+  if (!targetAttr)
+    targetAttr = ExecutableTargetAttr::lookup(operation);
   auto materializeEncodingFn = getMaterializeEncodingFn(targetAttr);
   if (!materializeEncodingFn) {
     return signalPassFailure();
@@ -269,8 +283,10 @@ void CPUMaterializeEncodingPass::runOnOperation() {
 void CPUMaterializeUpperBoundTileSizePass::runOnOperation() {
   MLIRContext *context = &getContext();
   auto operation = getOperation();
-  auto targetAttrs =
-      IREE::HAL::DeviceTargetAttr::lookupExecutableTargets(operation);
+  if (targetAttrs.empty()) {
+    targetAttrs =
+        IREE::HAL::DeviceTargetAttr::lookupExecutableTargets(operation);
+  }
   RewritePatternSet patterns(context);
   MaterializeEncodingFn materializeEncodingFn =
       getUpperBoundMaterializeEncodingFn(targetAttrs);
@@ -290,11 +306,20 @@ std::unique_ptr<OperationPass<func::FuncOp>>
 createCPUMaterializeEncodingPass() {
   return std::make_unique<CPUMaterializeEncodingPass>();
 }
+std::unique_ptr<OperationPass<func::FuncOp>>
+createCPUMaterializeEncodingPass(IREE::HAL::ExecutableTargetAttr targetAttr) {
+  return std::make_unique<CPUMaterializeEncodingPass>(targetAttr);
+}
 
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
 createCPUMaterializeUpperBoundTileSizePass() {
   return std::make_unique<CPUMaterializeUpperBoundTileSizePass>();
 }
+std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
+createCPUMaterializeUpperBoundTileSizePass(
+    ArrayRef<IREE::HAL::ExecutableTargetAttr> targetAttrs) {
+  return std::make_unique<CPUMaterializeUpperBoundTileSizePass>(targetAttrs);
+}
 
 } // namespace iree_compiler
 } // namespace mlir

compiler/src/iree/compiler/Codegen/Common/CPU/Passes.h

@@ -12,6 +12,7 @@
 #ifndef IREE_COMPILER_CODEGEN_COMMON_CPU_PASSES_H_
 #define IREE_COMPILER_CODEGEN_COMMON_CPU_PASSES_H_
 
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Pass/Pass.h"
 
@@ -23,6 +24,8 @@ namespace iree_compiler {
 ///   linalg_ext.unset_encoding -> tensor.unpack
 ///   linalg.matmul             -> linalg.mmt4d
 std::unique_ptr<OperationPass<func::FuncOp>> createCPUMaterializeEncodingPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createCPUMaterializeEncodingPass(IREE::HAL::ExecutableTargetAttr targetAttr);
 
 /// Like createLLVMCPUMaterializeEncodingPass, but specifically for
 /// linalg_ext.upper_bound_tile_size, converting it to constants.
@@ -41,6 +44,9 @@ std::unique_ptr<OperationPass<func::FuncOp>> createCPUMaterializeEncodingPass();
 // as needed.
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
 createCPUMaterializeUpperBoundTileSizePass();
+std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
+createCPUMaterializeUpperBoundTileSizePass(
+    ArrayRef<IREE::HAL::ExecutableTargetAttr> targetAttrs);
 
 void registerCodegenCommonCPUPasses();
 

compiler/src/iree/compiler/Dialect/HAL/Target/LLVMCPU/BUILD.bazel

@@ -36,6 +36,7 @@ iree_compiler_cc_library(
         ":LinkerTool",
         ":StaticLibraryGenerator",
         "//compiler/src/iree/compiler/Codegen/Common",
+        "//compiler/src/iree/compiler/Codegen/Common/CPU:CommonCPUPasses",
         "//compiler/src/iree/compiler/Codegen/Dialect:IREECodegenDialect",
         "//compiler/src/iree/compiler/Codegen/LLVMCPU",
         "//compiler/src/iree/compiler/Codegen/Utils",

compiler/src/iree/compiler/Dialect/HAL/Target/LLVMCPU/CMakeLists.txt

@@ -48,6 +48,7 @@ iree_cc_library(
     MLIRTargetLLVMIRExport
     MLIRTransformDialect
     iree::compiler::Codegen::Common
+    iree::compiler::Codegen::Common::CPU::CommonCPUPasses
     iree::compiler::Codegen::Dialect::IREECodegenDialect
     iree::compiler::Codegen::LLVMCPU
     iree::compiler::Codegen::Utils

compiler/src/iree/compiler/Dialect/HAL/Target/LLVMCPU/LLVMCPUTarget.cpp

@@ -11,6 +11,7 @@
 
 #include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
 #include "iree-dialects/Dialect/LinalgTransform/LinalgTransformOps.h"
+#include "iree/compiler/Codegen/Common/CPU/Passes.h"
 #include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/LLVMCPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
@@ -218,6 +219,15 @@ class LLVMCPUTargetBackend final : public TargetBackend {
     buildLLVMCPUCodegenPassPipeline(passManager);
   }
 
+  void buildMaterializeEncodingsPassPipeline(
+      IREE::HAL::ExecutableTargetAttr executableTarget,
+      OpPassManager &passManager) override {
+    passManager.addNestedPass<func::FuncOp>(
+        createCPUMaterializeUpperBoundTileSizePass(executableTarget));
+    passManager.addNestedPass<func::FuncOp>(
+        createCPUMaterializeEncodingPass({executableTarget}));
+  }
+
   void buildLinkingPassPipeline(OpPassManager &passManager) override {
     buildLLVMCPULinkingPassPipeline(passManager);
   }

compiler/src/iree/compiler/Dialect/HAL/Target/TargetBackend.h

@@ -185,6 +185,10 @@ class TargetBackend {
   buildTranslationPassPipeline(IREE::HAL::ExecutableVariantOp variantOp,
                                OpPassManager &passManager) = 0;
 
+  virtual void buildMaterializeEncodingsPassPipeline(
+      IREE::HAL::ExecutableTargetAttr executableTarget,
+      OpPassManager &passManager) {}
+
   // Inserts passes used to link `hal.executable.variant` ops together.
   // The pass manager will be nested on the parent module of `hal.executable`
   // ops and the pipeline will need to find relevant variant ops itself.

compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD.bazel

@@ -25,6 +25,7 @@ iree_compiler_cc_library(
         "InlineDeviceSwitches.cpp",
         "LinkExecutables.cpp",
         "MaterializeDispatchInstrumentation.cpp",
+        "MaterializeHomogeneousEncodings.cpp",
         "MaterializeInterfaces.cpp",
         "MaterializeResourceCaches.cpp",
         "MemoizeDeviceQueries.cpp",

compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt

@@ -26,6 +26,7 @@ iree_cc_library(
     "InlineDeviceSwitches.cpp"
     "LinkExecutables.cpp"
     "MaterializeDispatchInstrumentation.cpp"
+    "MaterializeHomogeneousEncodings.cpp"
     "MaterializeInterfaces.cpp"
     "MaterializeResourceCaches.cpp"
     "MemoizeDeviceQueries.cpp"

compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeHomogeneousEncodings.cpp

@@ -0,0 +1,90 @@
+// 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/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/HAL/Target/TargetBackend.h"
+#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
+#include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace HAL {
+
+class MaterializeHomogeneousEncodingsPass
+    : public PassWrapper<MaterializeHomogeneousEncodingsPass,
+                         OperationPass<ModuleOp>> {
+public:
+  MaterializeHomogeneousEncodingsPass()
+      : targetRegistry(TargetBackendRegistry::getGlobal()) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<IREE::HAL::HALDialect>();
+  }
+
+  StringRef getArgument() const override {
+    return "iree-hal-materialize-homogeneous-encodings";
+  }
+
+  StringRef getDescription() const override {
+    return "Mateiralizes logical encodings to physical encodings if there is "
+           "a single device target.";
+  }
+
+  void runOnOperation() override {
+    auto moduleOp = getOperation();
+    auto targetsAttr = moduleOp->getAttrOfType<ArrayAttr>("hal.device.targets");
+    if (!targetsAttr || targetsAttr.size() != 1) {
+      return;
+    }
+    auto deviceTarget = cast<IREE::HAL::DeviceTargetAttr>(targetsAttr[0]);
+    SmallVector<ExecutableTargetAttr, 4> executableTargets =
+        deviceTarget.getExecutableTargets();
+    if (executableTargets.size() != 1) {
+      return;
+    }
+    auto executableTarget = executableTargets[0];
+    OpPassManager passManager(moduleOp.getOperationName());
+    auto targetBackend =
+        targetRegistry.getTargetBackend(executableTarget.getBackend());
+    if (!targetBackend) {
+      moduleOp.emitError() << "unregistered target backend '"
+                           << executableTarget.getBackend() << "'";
+      return;
+    }
+    targetBackend->buildMaterializeEncodingsPassPipeline(executableTarget,
+                                                         passManager);
+    if (failed(runPipeline(passManager, moduleOp))) {
+      return signalPassFailure();
+    }
+  }
+
+private:
+  const TargetBackendRegistry &targetRegistry;
+};
+
+std::unique_ptr<OperationPass<ModuleOp>>
+createMaterializeHomogeneousEncodingsPass() {
+  return std::make_unique<MaterializeHomogeneousEncodingsPass>();
+}
+
+static PassRegistration<MaterializeHomogeneousEncodingsPass> pass([] {
+  return std::make_unique<MaterializeHomogeneousEncodingsPass>();
+});
+
+} // namespace HAL
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir

compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp

@@ -227,8 +227,9 @@ void buildHALTransformPassPipeline(OpPassManager &passManager,
     // Executable translation
     //----------------------------------------------------------------------------
 
-    FunctionLikeNest(passManager)
-        .addPass(createCPUMaterializeUpperBoundTileSizePass);
+    FunctionLikeNest(passManager).addPass([]() {
+      return createCPUMaterializeUpperBoundTileSizePass();
+    });
 
     // Preprocess executables using an external tool. The tool may mutate one or
     // more variants and even insert or remove variants.

compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h

@@ -106,6 +106,9 @@ std::unique_ptr<OperationPass<mlir::ModuleOp>> createMemoizeDeviceQueriesPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>>
 createMaterializeInterfacesPass();
 
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createMaterializeHomogeneousEncodingsPass();
+
 // Dumps individual hal.executable source listings to |path|.
 std::unique_ptr<OperationPass<mlir::ModuleOp>>
 createDumpExecutableSourcesPass(StringRef path);
@@ -213,6 +216,7 @@ inline void registerHALPasses() {
   createLinkExecutablesPass(TargetBackendRegistry::getGlobal());
   createLinkTargetExecutablesPass(TargetBackendRegistry::getGlobal(), "");
   createMaterializeDispatchInstrumentationPass(0);
+  createMaterializeHomogeneousEncodingsPass();
   createMaterializeInterfacesPass();
   createMaterializeResourceCachesPass(targetOptions);
   createMemoizeDeviceQueriesPass();

compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD.bazel

@@ -25,6 +25,7 @@ iree_lit_test_suite(
             "fixup_legacy_sync.mlir",
             "inline_device_switches.mlir",
             "materialize_dispatch_instrumentation.mlir",
+            "materialize_homogeneous_encodings.mlir",
             "materialize_interfaces.mlir",
             "materialize_resource_caches.mlir",
             "memoize_device_queries.mlir",

compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt

@@ -23,6 +23,7 @@ iree_lit_test_suite(
     "fixup_legacy_sync.mlir"
     "inline_device_switches.mlir"
     "materialize_dispatch_instrumentation.mlir"
+    "materialize_homogeneous_encodings.mlir"
     "materialize_interfaces.mlir"
     "materialize_resource_caches.mlir"
     "memoize_device_queries.mlir"