Add support for executable duplication in encoding specialization pass.

compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.cpp

@@ -7,6 +7,7 @@
 #include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "iree/compiler/Utils/StringUtils.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
@@ -893,6 +894,30 @@ bool DeviceAffinityAttr::isExecutableWith(
   return false;
 }
 
+bool DeviceAffinityAttr::isTranslatableWith(
+    ModuleOp moduleOp, IREE::Stream::AffinityAttr other) const {
+  if (!other)
+    return true;
+
+  auto otherAffinityAttr = llvm::dyn_cast_if_present<DeviceAffinityAttr>(other);
+  if (!otherAffinityAttr)
+    return false;
+
+  SymbolTable symbolTable(moduleOp);
+  auto getExecutableTargets = [&](SymbolRefAttr symbol) {
+    auto globalOp = llvm::cast<IREE::Util::GlobalOp>(
+        symbolTable.lookupSymbolIn(moduleOp, symbol));
+    return globalOp.getInitialValueAttr();
+  };
+  auto device = llvm::dyn_cast<IREE::HAL::DeviceTargetAttr>(
+      getExecutableTargets(getDevice()));
+  auto otherDevice = llvm::dyn_cast<IREE::HAL::DeviceTargetAttr>(
+      getExecutableTargets(otherAffinityAttr.getDevice()));
+  if (!device || !otherDevice)
+    return false;
+  return device.getExecutableTargets() == otherDevice.getExecutableTargets();
+}
+
 IREE::Stream::AffinityAttr
 DeviceAffinityAttr::joinOR(IREE::Stream::AffinityAttr other) const {
   if (!other)

compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.td

@@ -853,6 +853,7 @@ def HAL_DeviceSelectAttr : AttrDef<HAL_Dialect, "DeviceSelect", [
 def HAL_DeviceAffinityAttr : AttrDef<HAL_Dialect, "DeviceAffinity", [
   DeclareAttrInterfaceMethods<Stream_AffinityAttr, [
     "isExecutableWith",
+    "isTranslatableWith",
     "joinOR",
     "joinAND",
   ]>,

compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td

@@ -59,6 +59,20 @@ def Stream_AffinityAttr : AttrInterface<"AffinityAttr"> {
         return IREE::Stream::AffinityAttr::areCompatible($_attr, other);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns true if it can share the same execution configuration (e.g.,
+        translation artifacts) with the `other` and vice versa.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isTranslatableWith",
+      /*args=*/(ins "ModuleOp":$moduleOp,
+                    "IREE::Stream::AffinityAttr":$other),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return IREE::Stream::AffinityAttr::areTranslationCompatible(moduleOp, $_attr, other);
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Returns an affinity describing the union with |other| constraints.
@@ -118,6 +132,11 @@ def Stream_AffinityAttr : AttrInterface<"AffinityAttr"> {
     // Returns true if |lhs| and |rhs| indicate that their operations can
     // execute together on the same execution queue.
     static bool canExecuteTogether(AffinityAttr lhs, AffinityAttr rhs);
+
+    // Returns true if |lhs| and |rhs| are translation compatible. E.g., they
+    // are compatible if they have the same executable targets.
+    static bool areTranslationCompatible(ModuleOp moduleOp, AffinityAttr lhs,
+                                         AffinityAttr rhs);
   }];
 }
 

compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp

@@ -393,6 +393,16 @@ bool AffinityAttr::canExecuteTogether(AffinityAttr lhs, AffinityAttr rhs) {
   return lhs.isExecutableWith(rhs);
 }
 
+// static
+bool AffinityAttr::areTranslationCompatible(ModuleOp moduleOp, AffinityAttr lhs,
+                                            AffinityAttr rhs) {
+  if (lhs == rhs)
+    return true;
+  if ((lhs && !rhs) || (rhs && !lhs))
+    return true;
+  return lhs.isTranslatableWith(moduleOp, rhs);
+}
+
 //===----------------------------------------------------------------------===//
 // #stream.partitioning_config
 //===----------------------------------------------------------------------===//