[Codegen] Add pass for specializing executable variants (#20771)

This allows for kernel specialization by way of annotating tilable ops
with the preferred iteration ranges and divisibilities to specialize
for. For example,

```
%35 = linalg.matmul ins(%30, %31 : tensor<?x?xf16>, tensor<?x?xf16>) outs(%34 : tensor<?x?xf32>)
  {iree_codegen.specialization_ranges = #util<int.assumption.multi_array[
    [<umin = 128, umax = 4096, udiv = 128>, <umin = 128, umax = 4096, udiv = 128>, <umin = 64, udiv = 64>],
    [<umin = 4096, udiv = 256>, <umin = 4096, udiv = 256>, <udiv = 64>]
]>} -> tensor<?x?xf32>
```

This matmul will be specialized into up to 3 variants, one default, and
two for each list of ranges. More details can be found in the code
documentation.

To specialize a dispatch, the pass queries the iteration domain of the
tilable op and looks for a producer `util.int.assume` op. The function
is then cloned and the assume is updated to match the requested ranges.

This is just one way to implement kernel specialization, in the future
we can implement many more options (e.g. based on estimated TOPs or
smallest individual dimension) and that's without considering what
device we're targeting. The implementation here supports many common
cases of minimum/maximum sizes and alignments. Learnings here will
translate to future improvements for other strategies.

This pass runs first thing in the configuration pipelines because it
needs to apply to the variant, and should run before passes like
BlockDynamicDims which takes advantage of `util.assume.int` information
that this pass refines.

Author: qedawkins

compiler/src/iree/compiler/Codegen/Common/BUILD.bazel

@@ -156,6 +156,7 @@ iree_compiler_cc_library(
         "RemoveSingleIterationLoop.cpp",
         "ReplaceSlowMinMaxOps.cpp",
         "ResolveSwizzleHints.cpp",
+        "SpecializeExports.cpp",
         "SplitFullPartialTransferPass.cpp",
         "StripCompilationInfoPass.cpp",
         "TensorDynamicDimAnalysis.cpp",

compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt

@@ -145,6 +145,7 @@ iree_cc_library(
     "RemoveSingleIterationLoop.cpp"
     "ReplaceSlowMinMaxOps.cpp"
     "ResolveSwizzleHints.cpp"
+    "SpecializeExports.cpp"
     "SplitFullPartialTransferPass.cpp"
     "StripCompilationInfoPass.cpp"
     "TensorDynamicDimAnalysis.cpp"

compiler/src/iree/compiler/Codegen/Common/Passes.td

@@ -658,6 +658,11 @@ def ResolveSwizzleHintsPass :
   ];
 }
 
+def SpecializeExportsPass :
+    Pass<"iree-codegen-specialize-exports", "IREE::HAL::ExecutableVariantOp">{
+   let summary = "Specializes exported functions based on annotated ranges";
+}
+
 def StripCompilationInfoPass :
     Pass<"iree-codegen-strip-compilation-info", "">{
    let summary = "Remove all the the lowering configuration and translation info attributes.";