[Codegen] Add vector size inference for ukernel operations. (iree-org#22440)

## Summary

This change enables vector size inference for
`iree_codegen.ukernel.generic` operations, allowing downstream
operations like `linalg.unpack` to properly vectorize when consuming
ukernel outputs.

## Problem

The vector size inference mechanism in `inferSizesFromIR` did not handle
`UKernelGenericOp`, causing operations that consume ukernel outputs to
fail vectorization. This resulted in scalar code generation instead of
vectorized operations, leading to suboptimal performance.

Example: An `linalg.unpack` operation following a `iree_uk_mmt4d`
ukernel would use vector sizes of `[1, 1]` instead of the correct `[16,
16]` based on the ukernel's output shape.

## Solution

Added a case in the `inferSizesFromIR` TypeSwitch to handle
`UKernelGenericOp`:
- Infer vector sizes directly from the result tensor type shape
- Handle dynamic dimensions using `ValueBoundsConstraintSet`
- Return `VectorizationTileSizes` with both `vectorSizes` and
`destShape` populated

## Changes

**compiler/src/iree/compiler/Codegen/Utils/Utils.cpp:**
- Added `#include
"iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"`
- Added standalone `inferSizesFromIR(UKernelGenericOp, OpResult)`
function following the pattern of other operations
- Updated TypeSwitch in `inferSizesFromIR(Value val)` to call the new
function
- Used `llvm::enumerate` with `static_cast<unsigned>` for dimension
indexing

**compiler/src/iree/compiler/Codegen/Utils/Utils.h:**
- Added `#include
"iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"` in proper
alphabetical order
- Added public function declaration for
`inferSizesFromIR(UKernelGenericOp, OpResult)`


**compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization.mlir:**
- Added test `@ukernel_unpack_infer_vector_sizes` verifying
vectorization of unpack after ukernel

## Testing

The test verifies that:
1. A ukernel operation produces a `tensor<1x1x16x16xf32>` 
2. The following unpack operation correctly infers vector sizes and
vectorizes to `vector<16x16xf32>`
3. The vectorization uses `vector.transfer_read`, `vector.shape_cast`,
and `vector.transfer_write` operations

Run test with:
```bash
iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-generic-vectorization{enable-vector-masking=true}))" test.mlir
```

This change bridges the gap in the vector size inference mechanism,
enabling proper vectorization for any operations that consume ukernel
outputs, not just unpack operations.

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: hanhanW <[email protected]>

Author: 2 people

compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization.mlir

@@ -811,3 +811,28 @@ func.func @negative_no_vectorize_large_vector(%arg0 : tensor<1x9007199254740991x
 // CHECK-MASK:           } -> tensor<1x9007199254740991xf32>
 // CHECK-MASK:           return %[[VAL_2]] : tensor<1x9007199254740991xf32>
 // CHECK-MASK:         }
+
+// -----
+
+// Test that unpack operations following ukernel operations can properly infer
+// vector sizes from the ukernel output shape.
+
+func.func @ukernel_unpack_infer_vector_sizes(%lhs: tensor<1x8x16x1xf32>, %rhs: tensor<1x8x16x1xf32>, %dest: tensor<16x16xf32>) -> tensor<16x16xf32> {
+  %init = tensor.empty() : tensor<1x1x16x16xf32>
+  %ukernel = iree_codegen.ukernel.generic "foo"
+    ins(%lhs, %rhs : tensor<1x8x16x1xf32>, tensor<1x8x16x1xf32>)
+    outs(%init : tensor<1x1x16x16xf32>)
+    -> tensor<1x1x16x16xf32>
+  %unpack = linalg.unpack %ukernel
+    outer_dims_perm = [0, 1]
+    inner_dims_pos = [0, 1]
+    inner_tiles = [16, 16]
+    into %dest
+    : tensor<1x1x16x16xf32> -> tensor<16x16xf32>
+  return %unpack : tensor<16x16xf32>
+}
+// CHECK-MASK-LABEL: func.func @ukernel_unpack_infer_vector_sizes
+// CHECK-MASK:         %[[UKERNEL:.*]] = iree_codegen.ukernel.generic "foo"
+// CHECK-MASK:         %[[READ:.*]] = vector.transfer_read %[[UKERNEL]]{{.*}} : tensor<1x1x16x16xf32>, vector<1x1x16x16xf32>
+// CHECK-MASK:         %[[CAST:.*]] = vector.shape_cast %[[READ]] : vector<1x1x16x16xf32> to vector<16x16xf32>
+// CHECK-MASK:         vector.transfer_write %[[CAST]]{{.*}} : vector<16x16xf32>, tensor<16x16xf32>

compiler/src/iree/compiler/Codegen/Utils/Utils.cpp

@@ -7,6 +7,7 @@
 #include "iree/compiler/Codegen/Utils/Utils.h"
 
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/Interfaces/ProcessorOpInterfaces.h"
 #include "iree/compiler/Codegen/Interfaces/UKernelOpInterface.h"
@@ -1912,6 +1913,44 @@ std::optional<VectorizationTileSizes> inferSizesFromIR(linalg::UnPackOp op) {
   return result;
 }
 
+std::optional<VectorizationTileSizes>
+inferSizesFromIR(IREE::Codegen::UKernelGenericOp ukernelOp, OpResult opResult) {
+  LDBG() << "Inferring dest sizes for: " << ukernelOp;
+  auto resultType = dyn_cast<RankedTensorType>(opResult.getType());
+  if (!resultType) {
+    LDBG()
+        << "failed to infer sizes because result type is not a ranked tensor";
+    return std::nullopt;
+  }
+
+  VectorizationTileSizes result;
+  for (auto [idx, dim] : llvm::enumerate(resultType.getShape())) {
+    if (ShapedType::isDynamic(dim)) {
+      FailureOr<int64_t> maybeDimBound =
+          ValueBoundsConstraintSet::computeConstantBound(
+              presburger::BoundType::UB, {opResult, static_cast<unsigned>(idx)},
+              /*stopCondition=*/nullptr, /*closedUB=*/true);
+      if (failed(maybeDimBound)) {
+        LDBG() << "failed to infer bounds for dynamic dim";
+        return std::nullopt;
+      }
+      result.vectorSizes.push_back(maybeDimBound.value());
+    } else {
+      result.vectorSizes.push_back(dim);
+    }
+  }
+  result.destShape = result.vectorSizes;
+
+  LLVM_DEBUG({
+    LDBG() << "Inferred vector sizes:";
+    for (auto [idx, val] : llvm::enumerate(result.vectorSizes)) {
+      LDBG() << "Dim #" << idx << ": " << val;
+    }
+  });
+
+  return result;
+}
+
 std::optional<VectorizationTileSizes> static inferSizesFromMixedSizes(
     SmallVector<OpFoldResult> shape) {
   VectorizationTileSizes result;
@@ -1948,6 +1987,8 @@ std::optional<VectorizationTileSizes> inferSizesFromIR(Value val) {
         // the values.
         result = inferSizesFromMixedSizes(op.getMixedSizes());
       })
+      .Case<IREE::Codegen::UKernelGenericOp>(
+          [&](auto op) { result = inferSizesFromIR(op, cast<OpResult>(val)); })
       .Default([&](Operation *) {});
 
   return result;

compiler/src/iree/compiler/Codegen/Utils/Utils.h

@@ -9,6 +9,7 @@
 
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
@@ -365,6 +366,12 @@ inferSizesFromIR(linalg::LinalgOp linalgOp, std::optional<OpResult> opResult);
 std::optional<VectorizationTileSizes> inferSizesFromIR(scf::ForOp forOp,
                                                        OpResult opResult);
 
+/// Returns the result sizes and vector input sizes of the ukernel.generic op.
+/// The inferred bounding size is returned if it is dynamic shape. Returns
+/// std::nullopt if the shape inference failed.
+std::optional<VectorizationTileSizes>
+inferSizesFromIR(IREE::Codegen::UKernelGenericOp ukernelOp, OpResult opResult);
+
 /// Returns the underlying index if the given value is a constant index.
 std::optional<int64_t> getConstantIndex(Value value);