[Codegen][VectorDistribute] Add pattern to distribute poison (#21573)

newling · web-flow · commit 9323cd236fb9 · 2025-08-06T20:56:06.000Z
Essentially the same as the pattern to distribute constants.

Also removes some unused headers from touched file.

---------

Signed-off-by: James Newling &lt;james.newling@gmail.com&gt;
diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
@@ -85,6 +85,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:UBDialect",
         "@llvm-project//mlir:VectorDialect",
     ],
 )
diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
@@ -60,6 +60,7 @@ iree_cc_library(
     MLIRAnalysis
     MLIRIR
     MLIRSCFDialect
+    MLIRUBDialect
     MLIRVectorDialect
     iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
   PUBLIC
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
@@ -48,6 +48,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:TransformUtils",
         "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:UBDialect",
     ],
 )
 
@@ -167,6 +168,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:TilingInterface",
         "@llvm-project//mlir:TransformUtils",
         "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:UBDialect",
         "@llvm-project//mlir:ValueBoundsOpInterface",
         "@llvm-project//mlir:VectorDialect",
         "@llvm-project//mlir:VectorToSCF",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
@@ -33,6 +33,7 @@ iree_cc_library(
     MLIRPass
     MLIRTransformUtils
     MLIRTransforms
+    MLIRUBDialect
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Utils
@@ -135,6 +136,7 @@ iree_cc_library(
     MLIRTilingInterface
     MLIRTransformUtils
     MLIRTransforms
+    MLIRUBDialect
     MLIRValueBoundsOpInterface
     MLIRVectorDialect
     MLIRVectorToSCF
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUDistributionPatterns.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUDistributionPatterns.cpp
@@ -4,19 +4,16 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include <numeric>
 #include "iree/compiler/Codegen/Common/GPU/GPUPatterns.h"
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
-#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
-#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
-#include "iree/compiler/Codegen/Utils/GPUUtils.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 
@@ -57,6 +54,30 @@ struct DistributeConstants final : OpDistributionPattern<arith::ConstantOp> {
   }
 };
 
+struct DistributePoison final : OpDistributionPattern<ub::PoisonOp> {
+  using OpDistributionPattern::OpDistributionPattern;
+
+  LogicalResult matchAndRewrite(ub::PoisonOp poisonOp,
+                                DistributionSignature &signature,
+                                PatternRewriter &rewriter) const override {
+
+    auto poisonVal = dyn_cast<VectorValue>(poisonOp.getResult());
+    if (!poisonVal)
+      return failure();
+
+    SmallVector<int64_t> distributedShape =
+        signature[poisonVal].getDistributedShape();
+
+    Type elementType = poisonVal.getType().getElementType();
+    auto vectorType = VectorType::get(distributedShape, elementType);
+    auto distributedOp =
+        ub::PoisonOp::create(rewriter, poisonVal.getLoc(), vectorType);
+    replaceOpWithDistributedValues(rewriter, poisonOp,
+                                   distributedOp->getResult(0));
+    return success();
+  }
+};
+
 struct DistributeElementwise final
     : OpTraitDistributionPattern<OpTrait::Elementwise> {
   using OpTraitDistributionPattern::OpTraitDistributionPattern;
@@ -336,8 +357,8 @@ struct DistributeTrivialExtract final
 } // namespace
 
 void populateGPUDistributionPatterns(RewritePatternSet &patterns) {
-  patterns.add<DistributeConstants, DistributeScfFor, DistributeTrivialExtract>(
-      patterns.getContext());
+  patterns.add<DistributeConstants, DistributePoison, DistributeScfFor,
+               DistributeTrivialExtract>(patterns.getContext());
   // Elementwise patterns.
   patterns.add<DistributeElementwise>(patterns.getContext());
   patterns.add<DistributeTrivialLayoutConversions>(patterns.getContext());
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir
@@ -1,14 +1,13 @@
-// RUN: iree-opt --iree-transform-dialect-interpreter --split-input-file --canonicalize --cse %s | FileCheck %s
+// RUN: iree-opt --iree-transform-dialect-interpreter --split-input-file %s | FileCheck %s
 
 #nested = #iree_vector_ext.nested_layout<
-  subgroup_tile = [2, 1, 1],
-  batch_tile    = [8, 2, 4],
-  outer_tile        = [1, 4, 4],
-  thread_tile       = [8, 2, 4],
+  subgroup_tile    = [2, 1, 1],
+  batch_tile       = [8, 2, 4],
+  outer_tile       = [1, 4, 4],
+  thread_tile      = [8, 2, 4],
   element_tile     = [1, 8, 2],
-
-  subgroup_strides        = [1, 1, 1],
-  thread_strides          = [1, 8, 16]
+  subgroup_strides = [1, 1, 1],
+  thread_strides   = [1, 8, 16]
 >
 
 // CHECK-LABEL: @distribute_elementwise_nested_layout_f16
@@ -28,15 +27,22 @@ func.func @distribute_elementwise_nested_layout_f16(%a: vector<128x128x128xf16>,
   return %d : vector<128x128x128xf16>
 }
 
-#layout = #iree_vector_ext.nested_layout<
-  subgroup_tile = [1, 1],
-  batch_tile = [1, 1],
-  outer_tile = [1, 1],
-  thread_tile = [1, 1],
-  element_tile = [16, 16],
+// CHECK-LABEL: @distribute_poison
+func.func @distribute_poison() -> vector<128x128x128xf16> {
+  // CHECK: ub.poison : vector<8x2x4x1x4x4x1x8x2xf16>
+  %root = ub.poison : vector<128x128x128xf16>
+  %rootl = iree_vector_ext.to_layout %root to layout(#nested) : vector<128x128x128xf16>
+  return %rootl: vector<128x128x128xf16>
+}
 
+#layout = #iree_vector_ext.nested_layout<
+  subgroup_tile    = [1, 1],
+  batch_tile       = [1, 1],
+  outer_tile       = [1, 1],
+  thread_tile      = [1, 1],
+  element_tile     = [16, 16],
   subgroup_strides = [1, 1],
-  thread_strides = [1, 1]
+  thread_strides   = [1, 1]
 >
 
 // CHECK-LABEL: @distribute_scf_for
@@ -63,14 +69,13 @@ func.func @distribute_scf_for(%a: vector<16x16xi32>, %b: vector<16x16xi32>) -> v
 }
 
 #layout_0d = #iree_vector_ext.nested_layout<
-  subgroup_tile = [],
-  batch_tile = [],
-  outer_tile = [],
-  thread_tile = [],
-  element_tile = [],
-
+  subgroup_tile    = [],
+  batch_tile       = [],
+  outer_tile       = [],
+  thread_tile      = [],
+  element_tile     = [],
   subgroup_strides = [],
-  thread_strides = []
+  thread_strides   = []
 >
 
 // CHECK-LABEL: @distribute_scf_for_0d

Original file line number	Diff line number	Diff line change
`@@ -85,6 +85,7 @@ iree_compiler_cc_library(`
`85`	`85`	`"@llvm-project//mlir:DialectUtils",`
`86`	`86`	`"@llvm-project//mlir:IR",`
`87`	`87`	`"@llvm-project//mlir:SCFDialect",`
	`88`	`+ "@llvm-project//mlir:UBDialect",`
`88`	`89`	`"@llvm-project//mlir:VectorDialect",`
`89`	`90`	`],`
`90`	`91`	`)`