[llvm] [mlir] [mlir][dataflow] Add uniformity analysis (PR #220740)
Alessandro Potenza via llvm-commits
llvm-commits at lists.llvm.org
Fri Sep 4 23:39:35 PDT 2026
https://github.com/alepot55 updated https://github.com/llvm/llvm-project/pull/220740
>From 2c938c225c0ae94453fd2c120e7c9db764c1b9e2 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Wed, 2 Sep 2026 22:20:55 +0200
Subject: [PATCH 1/3] [mlir][dataflow] Make visitRegionSuccessors overridable
The method is virtual and its comment says it may be overridden to control
precisely how the region successors of a branch are visited, but it is declared
in the private section, so an override cannot delegate to the default
implementation for the cases it does not want to handle itself.
Move the declaration to the protected section, next to the other hooks a
subclass is meant to override. NFC: only the access level changes. The twin in
AbstractSparseBackwardDataFlowAnalysis is left alone because it is not virtual,
so making it protected would not let an override delegate to it.
---
.../mlir/Analysis/DataFlow/SparseAnalysis.h | 28 +++++++++----------
1 file changed, 14 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
index df50d8d193aeb..fa313ae7a4440 100644
--- a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
@@ -257,6 +257,20 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
visitCallableOperation(CallableOpInterface callable,
ArrayRef<AbstractSparseLattice *> argLattices);
+ /// Visit a program point `point` with predecessors within a region branch
+ /// operation `branch`, which can either be the entry block of one of the
+ /// regions or the parent operation itself, and set either the argument or
+ /// parent result lattices.
+ /// This method can be overridden to control precisely how the region
+ /// successors of `branch` are visited. For example in order to precisely
+ /// control the order in which predecessor operand lattices are propagated
+ /// from. An override is responsible for visiting all the known predecessors
+ /// and propagating therefrom.
+ virtual void
+ visitRegionSuccessors(ProgramPoint *point, RegionBranchOpInterface branch,
+ RegionSuccessor successor,
+ ArrayRef<AbstractSparseLattice *> lattices);
+
private:
/// Recursively initialize the analysis on nested operations and blocks.
LogicalResult initializeRecursively(Operation *op);
@@ -272,20 +286,6 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
/// region terminators or callable callsites. Otherwise, the values are
/// determined from block predecessors.
void visitBlock(Block *block);
-
- /// Visit a program point `point` with predecessors within a region branch
- /// operation `branch`, which can either be the entry block of one of the
- /// regions or the parent operation itself, and set either the argument or
- /// parent result lattices.
- /// This method can be overridden to control precisely how the region
- /// successors of `branch` are visited. For example in order to precisely
- /// control the order in which predecessor operand lattices are propagated
- /// from. An override is responsible for visiting all the known predecessors
- /// and propagating therefrom.
- virtual void
- visitRegionSuccessors(ProgramPoint *point, RegionBranchOpInterface branch,
- RegionSuccessor successor,
- ArrayRef<AbstractSparseLattice *> lattices);
};
//===----------------------------------------------------------------------===//
>From e8c01e580bce0b981d26b571347a70a876a33775 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Thu, 3 Sep 2026 07:30:16 +0200
Subject: [PATCH 2/3] [mlir] Add InferUniformityOpInterface
The interface through which an operation describes the uniformity of the values it defines on a SIMT target, and the lattice value that the analysis consuming it will use: a scope among divergent, subgroup, workgroup, cluster and uniform, or uninitialized, whose join is the narrower scope.
---
mlir/include/mlir/Interfaces/CMakeLists.txt | 1 +
.../Interfaces/InferUniformityOpInterface.h | 120 ++++++++++++++++++
.../Interfaces/InferUniformityOpInterface.td | 76 +++++++++++
mlir/lib/Interfaces/CMakeLists.txt | 2 +
.../Interfaces/InferUniformityOpInterface.cpp | 52 ++++++++
mlir/unittests/Interfaces/CMakeLists.txt | 2 +
.../InferUniformityOpInterfaceTest.cpp | 73 +++++++++++
.../llvm-project-overlay/mlir/BUILD.bazel | 30 +++++
.../mlir/unittests/BUILD.bazel | 1 +
9 files changed, 357 insertions(+)
create mode 100644 mlir/include/mlir/Interfaces/InferUniformityOpInterface.h
create mode 100644 mlir/include/mlir/Interfaces/InferUniformityOpInterface.td
create mode 100644 mlir/lib/Interfaces/InferUniformityOpInterface.cpp
create mode 100644 mlir/unittests/Interfaces/InferUniformityOpInterfaceTest.cpp
diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt
index 6461c68423c73..463250a56c136 100644
--- a/mlir/include/mlir/Interfaces/CMakeLists.txt
+++ b/mlir/include/mlir/Interfaces/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_interface(InferIntDivisibilityOpInterface)
add_mlir_interface(InferIntRangeInterface)
add_mlir_interface(InferStridedMetadataInterface)
add_mlir_interface(InferTypeOpInterface)
+add_mlir_interface(InferUniformityOpInterface)
add_mlir_interface(LoopLikeInterface)
add_mlir_interface(MemOpInterfaces)
add_mlir_interface(ParallelCombiningOpInterface)
diff --git a/mlir/include/mlir/Interfaces/InferUniformityOpInterface.h b/mlir/include/mlir/Interfaces/InferUniformityOpInterface.h
new file mode 100644
index 0000000000000..b575313d55103
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/InferUniformityOpInterface.h
@@ -0,0 +1,120 @@
+//===- InferUniformityOpInterface.h - Uniformity ----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the definitions of the uniformity inference interface
+// defined in `InferUniformityOpInterface.td`, and the lattice value it works
+// with.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_INFERUNIFORMITYOPINTERFACE_H
+#define MLIR_INTERFACES_INFERUNIFORMITYOPINTERFACE_H
+
+#include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/ArrayRef.h"
+#include <optional>
+
+namespace mlir {
+
+/// The widest group of threads of a SIMT launch within which an SSA value is
+/// known to be the same for every thread, from the narrowest to the widest.
+///
+/// The groups nest: a thread belongs to a subgroup (warp, wave), a subgroup to
+/// a workgroup (thread block), a workgroup to a cluster, and every cluster to
+/// the launch. A value that is the same within a group is also the same within
+/// every narrower group, so the scopes are totally ordered and the meet of two
+/// scopes is the narrower one.
+enum class UniformityScope : uint8_t {
+ /// Two threads of the same subgroup may observe different values.
+ Divergent = 0,
+ /// The same for every thread of a subgroup.
+ Subgroup,
+ /// The same for every thread of a workgroup.
+ Workgroup,
+ /// The same for every thread of a cluster of workgroups.
+ Cluster,
+ /// The same for every thread of the launch.
+ Uniform,
+};
+
+/// Returns the narrower of two scopes.
+inline UniformityScope meet(UniformityScope lhs, UniformityScope rhs) {
+ return lhs < rhs ? lhs : rhs;
+}
+
+/// Returns the name of a scope as written in IR and diagnostics.
+StringRef stringifyUniformityScope(UniformityScope scope);
+
+/// Returns the scope with the given name, or nullopt if there is none.
+std::optional<UniformityScope> symbolizeUniformityScope(StringRef name);
+
+raw_ostream &operator<<(raw_ostream &os, UniformityScope scope);
+
+/// The uniformity lattice value of an SSA value: a scope, or uninitialized
+/// while the analysis has not reached the value yet.
+class Uniformity {
+public:
+ Uniformity() = default;
+ Uniformity(UniformityScope scope) : scope(scope) {}
+
+ static Uniformity getUniform() {
+ return Uniformity(UniformityScope::Uniform);
+ }
+ static Uniformity getDivergent() {
+ return Uniformity(UniformityScope::Divergent);
+ }
+
+ bool isUninitialized() const { return !scope.has_value(); }
+
+ UniformityScope getScope() const {
+ assert(!isUninitialized() && "querying an uninitialized uniformity");
+ return *scope;
+ }
+
+ bool operator==(const Uniformity &rhs) const { return scope == rhs.scope; }
+ bool operator!=(const Uniformity &rhs) const { return !(*this == rhs); }
+
+ /// The join of two lattice values: uninitialized is the neutral element,
+ /// otherwise the narrower scope. The pessimistic fixpoint is `Divergent`.
+ static Uniformity join(const Uniformity &lhs, const Uniformity &rhs) {
+ if (lhs.isUninitialized())
+ return rhs;
+ if (rhs.isUninitialized())
+ return lhs;
+ return Uniformity(meet(*lhs.scope, *rhs.scope));
+ }
+
+ /// The join of a list of values; uninitialized if the list is empty or all
+ /// of its elements are uninitialized.
+ static Uniformity join(ArrayRef<Uniformity> values) {
+ Uniformity result;
+ for (const Uniformity &value : values)
+ result = join(result, value);
+ return result;
+ }
+
+ void print(raw_ostream &os) const;
+
+private:
+ std::optional<UniformityScope> scope;
+};
+
+inline raw_ostream &operator<<(raw_ostream &os, const Uniformity &uniformity) {
+ uniformity.print(os);
+ return os;
+}
+
+/// The callback through which an operation implementing
+/// `InferUniformityOpInterface` reports the uniformity of a value it defines.
+using SetUniformityFn = llvm::function_ref<void(Value, UniformityScope)>;
+
+} // namespace mlir
+
+#include "mlir/Interfaces/InferUniformityOpInterface.h.inc"
+
+#endif // MLIR_INTERFACES_INFERUNIFORMITYOPINTERFACE_H
diff --git a/mlir/include/mlir/Interfaces/InferUniformityOpInterface.td b/mlir/include/mlir/Interfaces/InferUniformityOpInterface.td
new file mode 100644
index 0000000000000..e6da54601f174
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/InferUniformityOpInterface.td
@@ -0,0 +1,76 @@
+//===- InferUniformityOpInterface.td - Uniformity ----------*- tablegen -*-===//
+//
+// Part of the LLVM Project, 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface through which operations participate in uniformity
+// (divergence) analysis on SIMT targets.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_INFERUNIFORMITYOPINTERFACE
+#define MLIR_INTERFACES_INFERUNIFORMITYOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def InferUniformityOpInterface : OpInterface<"InferUniformityOpInterface"> {
+ let description = [{
+ Allows operations to participate in uniformity analysis, which determines
+ for each SSA value the widest group of threads (subgroup, workgroup,
+ cluster, the whole launch) within which every thread observes the same
+ value.
+
+ Operations that read thread identity (`gpu.thread_id`), communicate across
+ lanes (`gpu.subgroup_broadcast`, `gpu.shuffle`) or define values for a
+ region executed by a subset of threads (`gpu.warp_execute_on_lane_0`)
+ implement this interface to describe how uniformity flows through them.
+ Operations that compute a function of their operands only need not
+ implement it: the analysis joins the uniformity of the operands.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ InterfaceMethod<[{
+ Infer the uniformity of the values defined by this operation, given the
+ uniformity of its operands, by calling `setUniformity` for each of them.
+
+ The values that may be set are the results of the operation and the
+ entry block arguments of its regions. A value that is not set keeps the
+ uniformity computed by the analysis: for a region entry block argument
+ or a result that region control flow forwards a value to, the join of
+ the forwarded values; for any other value, divergent.
+
+ For an operation implementing `RegionBranchOpInterface`, setting any
+ successor input of a region successor (a result, or an entry block
+ argument of a region) makes the operation authoritative for all
+ successor inputs of that successor: the analysis then does not join the
+ forwarded values into them. This is how an operation whose results are
+ not the values yielded from its regions, such as a broadcast, describes
+ itself.
+
+ Operands whose uniformity is not known yet are passed as uninitialized;
+ `Uniformity::join` treats them as neutral.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"inferUniformity",
+ /*args=*/(ins "::llvm::ArrayRef<::mlir::Uniformity>":$operandUniformity,
+ "::mlir::SetUniformityFn":$setUniformity)>,
+
+ InterfaceMethod<[{
+ Returns true if the regions of this operation are the entry of a SIMT
+ launch: the threads executing them are not the threads executing the
+ operation itself, so the control flow enclosing the operation does not
+ narrow the uniformity of their execution. `gpu.launch` returns true.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"isLaunchBoundary",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{ return false; }]>
+ ];
+}
+
+#endif // MLIR_INTERFACES_INFERUNIFORMITYOPINTERFACE
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index d20d290c45c01..163f66e40eacd 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -13,6 +13,7 @@ set(LLVM_OPTIONAL_SOURCES
InferIntRangeInterface.cpp
InferStridedMetadataInterface.cpp
InferTypeOpInterface.cpp
+ InferUniformityOpInterface.cpp
LoopLikeInterface.cpp
MemOpInterfaces.cpp
MemorySlotInterfaces.cpp
@@ -85,6 +86,7 @@ add_mlir_library(MLIRInferStridedMetadataInterface
)
add_mlir_interface_library(InferTypeOpInterface)
+add_mlir_interface_library(InferUniformityOpInterface)
add_mlir_library(MLIRLoopLikeInterface
LoopLikeInterface.cpp
diff --git a/mlir/lib/Interfaces/InferUniformityOpInterface.cpp b/mlir/lib/Interfaces/InferUniformityOpInterface.cpp
new file mode 100644
index 0000000000000..c2188c13b18f7
--- /dev/null
+++ b/mlir/lib/Interfaces/InferUniformityOpInterface.cpp
@@ -0,0 +1,52 @@
+//===- InferUniformityOpInterface.cpp - Uniformity inference --------------===//
+//
+// Part of the LLVM Project, 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 "mlir/Interfaces/InferUniformityOpInterface.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+#include "mlir/Interfaces/InferUniformityOpInterface.cpp.inc"
+
+StringRef mlir::stringifyUniformityScope(UniformityScope scope) {
+ switch (scope) {
+ case UniformityScope::Divergent:
+ return "divergent";
+ case UniformityScope::Subgroup:
+ return "subgroup";
+ case UniformityScope::Workgroup:
+ return "workgroup";
+ case UniformityScope::Cluster:
+ return "cluster";
+ case UniformityScope::Uniform:
+ return "uniform";
+ }
+ llvm_unreachable("unknown uniformity scope");
+}
+
+std::optional<UniformityScope> mlir::symbolizeUniformityScope(StringRef name) {
+ return llvm::StringSwitch<std::optional<UniformityScope>>(name)
+ .Case("divergent", UniformityScope::Divergent)
+ .Case("subgroup", UniformityScope::Subgroup)
+ .Case("workgroup", UniformityScope::Workgroup)
+ .Case("cluster", UniformityScope::Cluster)
+ .Case("uniform", UniformityScope::Uniform)
+ .Default(std::nullopt);
+}
+
+raw_ostream &mlir::operator<<(raw_ostream &os, UniformityScope scope) {
+ return os << stringifyUniformityScope(scope);
+}
+
+void Uniformity::print(raw_ostream &os) const {
+ if (isUninitialized())
+ os << "<uninitialized>";
+ else
+ os << *scope;
+}
diff --git a/mlir/unittests/Interfaces/CMakeLists.txt b/mlir/unittests/Interfaces/CMakeLists.txt
index 193edda6bf8d8..805084ef9b1eb 100644
--- a/mlir/unittests/Interfaces/CMakeLists.txt
+++ b/mlir/unittests/Interfaces/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_unittest(MLIRInterfacesTests
ControlFlowInterfacesTest.cpp
DataLayoutInterfacesTest.cpp
InferIntRangeInterfaceTest.cpp
+ InferUniformityOpInterfaceTest.cpp
MemorySlotUtilsTest.cpp
SideEffectInterfacesTest.cpp
InferTypeOpInterfaceTest.cpp
@@ -20,6 +21,7 @@ mlir_target_link_libraries(MLIRInterfacesTests
MLIRFuncDialect
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
+ MLIRInferUniformityOpInterface
MLIRMemorySlotUtils
MLIRParser
MLIRSideEffectInterfaces
diff --git a/mlir/unittests/Interfaces/InferUniformityOpInterfaceTest.cpp b/mlir/unittests/Interfaces/InferUniformityOpInterfaceTest.cpp
new file mode 100644
index 0000000000000..6fda0362ae18d
--- /dev/null
+++ b/mlir/unittests/Interfaces/InferUniformityOpInterfaceTest.cpp
@@ -0,0 +1,73 @@
+//===- InferUniformityOpInterfaceTest.cpp - Uniformity lattice tests ------===//
+//
+// Part of the LLVM Project, 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 "mlir/Interfaces/InferUniformityOpInterface.h"
+
+#include <gtest/gtest.h>
+
+using namespace mlir;
+
+TEST(Uniformity, ScopesNest) {
+ EXPECT_LT(UniformityScope::Divergent, UniformityScope::Subgroup);
+ EXPECT_LT(UniformityScope::Subgroup, UniformityScope::Workgroup);
+ EXPECT_LT(UniformityScope::Workgroup, UniformityScope::Cluster);
+ EXPECT_LT(UniformityScope::Cluster, UniformityScope::Uniform);
+ EXPECT_EQ(meet(UniformityScope::Workgroup, UniformityScope::Subgroup),
+ UniformityScope::Subgroup);
+ EXPECT_EQ(meet(UniformityScope::Subgroup, UniformityScope::Workgroup),
+ UniformityScope::Subgroup);
+ EXPECT_EQ(meet(UniformityScope::Uniform, UniformityScope::Uniform),
+ UniformityScope::Uniform);
+}
+
+TEST(Uniformity, JoinIsTheNarrowerScope) {
+ Uniformity workgroup(UniformityScope::Workgroup);
+ Uniformity cluster(UniformityScope::Cluster);
+ EXPECT_EQ(Uniformity::join(workgroup, cluster), workgroup);
+ EXPECT_EQ(Uniformity::join(cluster, workgroup), workgroup);
+ EXPECT_EQ(
+ Uniformity::join(Uniformity::getUniform(), Uniformity::getDivergent()),
+ Uniformity::getDivergent());
+ EXPECT_EQ(Uniformity::join(workgroup, workgroup), workgroup);
+}
+
+TEST(Uniformity, UninitializedIsNeutral) {
+ Uniformity none;
+ Uniformity subgroup(UniformityScope::Subgroup);
+ EXPECT_TRUE(none.isUninitialized());
+ EXPECT_FALSE(subgroup.isUninitialized());
+ EXPECT_EQ(Uniformity::join(none, subgroup), subgroup);
+ EXPECT_EQ(Uniformity::join(subgroup, none), subgroup);
+ EXPECT_TRUE(Uniformity::join(none, none).isUninitialized());
+}
+
+TEST(Uniformity, JoinOfList) {
+ EXPECT_TRUE(Uniformity::join(ArrayRef<Uniformity>()).isUninitialized());
+ Uniformity allNone[] = {Uniformity(), Uniformity()};
+ EXPECT_TRUE(Uniformity::join(allNone).isUninitialized());
+ Uniformity values[] = {Uniformity(), Uniformity::getUniform(),
+ Uniformity(UniformityScope::Cluster), Uniformity(),
+ Uniformity(UniformityScope::Workgroup)};
+ EXPECT_EQ(Uniformity::join(values).getScope(), UniformityScope::Workgroup);
+}
+
+TEST(Uniformity, Names) {
+ EXPECT_EQ(stringifyUniformityScope(UniformityScope::Divergent), "divergent");
+ EXPECT_EQ(stringifyUniformityScope(UniformityScope::Uniform), "uniform");
+ for (UniformityScope scope :
+ {UniformityScope::Divergent, UniformityScope::Subgroup,
+ UniformityScope::Workgroup, UniformityScope::Cluster,
+ UniformityScope::Uniform}) {
+ std::optional<UniformityScope> back =
+ symbolizeUniformityScope(stringifyUniformityScope(scope));
+ ASSERT_TRUE(back.has_value());
+ EXPECT_EQ(*back, scope);
+ }
+ EXPECT_FALSE(symbolizeUniformityScope("warp").has_value());
+ EXPECT_FALSE(symbolizeUniformityScope("Uniform").has_value());
+}
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 1b74d16736a47..2efe8697a6751 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -1505,6 +1505,13 @@ td_library(
deps = [":OpBaseTdFiles"],
)
+td_library(
+ name = "InferUniformityOpInterfaceTdFiles",
+ srcs = ["include/mlir/Interfaces/InferUniformityOpInterface.td"],
+ includes = ["include"],
+ deps = [":OpBaseTdFiles"],
+)
+
td_library(
name = "LoopLikeInterfaceTdFiles",
srcs = ["include/mlir/Interfaces/LoopLikeInterface.td"],
@@ -8064,6 +8071,29 @@ cc_library(
],
)
+gentbl_cc_library(
+ name = "InferUniformityOpInterfaceIncGen",
+ tbl_outs = {
+ "include/mlir/Interfaces/InferUniformityOpInterface.h.inc": ["-gen-op-interface-decls"],
+ "include/mlir/Interfaces/InferUniformityOpInterface.cpp.inc": ["-gen-op-interface-defs"],
+ },
+ tblgen = ":mlir-tblgen",
+ td_file = "include/mlir/Interfaces/InferUniformityOpInterface.td",
+ deps = [":InferUniformityOpInterfaceTdFiles"],
+)
+
+cc_library(
+ name = "InferUniformityOpInterface",
+ srcs = ["lib/Interfaces/InferUniformityOpInterface.cpp"],
+ hdrs = ["include/mlir/Interfaces/InferUniformityOpInterface.h"],
+ includes = ["include"],
+ deps = [
+ ":IR",
+ ":InferUniformityOpInterfaceIncGen",
+ "//llvm:Support",
+ ],
+)
+
td_library(
name = "DataLayoutInterfacesTdFiles",
srcs = ["include/mlir/Interfaces/DataLayoutInterfaces.td"],
diff --git a/utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel
index 675211bdf0cb7..4ca355b4dc6d9 100644
--- a/utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel
@@ -74,6 +74,7 @@ cc_test(
"//mlir:IR",
"//mlir:InferIntRangeInterface",
"//mlir:InferTypeOpInterface",
+ "//mlir:InferUniformityOpInterface",
"//mlir:MemorySlotUtils",
"//mlir:Parser",
"//mlir:SideEffectInterfaces",
>From 877a33ca8a0e21909227774c9d65db9f42184ad8 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Sat, 5 Sep 2026 08:39:12 +0200
Subject: [PATCH 3/3] [mlir][dataflow] Add uniformity analysis
A sparse forward dataflow analysis that computes, for each SSA value, the widest group of threads of a SIMT launch within which every thread observes the same value, and the widest group that executes an operation together. Operations describe themselves through InferUniformityOpInterface; region control flow needs no per-operation knowledge.
---
.../Analysis/DataFlow/UniformityAnalysis.h | 202 ++++++++
mlir/lib/Analysis/CMakeLists.txt | 2 +
.../Analysis/DataFlow/UniformityAnalysis.cpp | 448 ++++++++++++++++++
.../DataFlow/test-uniformity-analysis.mlir | 341 +++++++++++++
mlir/test/lib/Analysis/CMakeLists.txt | 1 +
.../DataFlow/TestUniformityAnalysis.cpp | 66 +++
mlir/test/lib/Dialect/Test/CMakeLists.txt | 1 +
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 122 +++++
mlir/test/lib/Dialect/Test/TestOps.h | 1 +
mlir/test/lib/Dialect/Test/TestOps.td | 103 ++++
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
.../llvm-project-overlay/mlir/BUILD.bazel | 2 +
.../mlir/test/BUILD.bazel | 2 +
13 files changed, 1293 insertions(+)
create mode 100644 mlir/include/mlir/Analysis/DataFlow/UniformityAnalysis.h
create mode 100644 mlir/lib/Analysis/DataFlow/UniformityAnalysis.cpp
create mode 100644 mlir/test/Analysis/DataFlow/test-uniformity-analysis.mlir
create mode 100644 mlir/test/lib/Analysis/DataFlow/TestUniformityAnalysis.cpp
diff --git a/mlir/include/mlir/Analysis/DataFlow/UniformityAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/UniformityAnalysis.h
new file mode 100644
index 0000000000000..a8a38259e5fb1
--- /dev/null
+++ b/mlir/include/mlir/Analysis/DataFlow/UniformityAnalysis.h
@@ -0,0 +1,202 @@
+//===- UniformityAnalysis.h - Uniformity analysis ---------------*- C++ -*-===//
+//
+// Part of the LLVM Project, 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the dataflow analysis that computes, for SIMT programs
+// expressed with structured control flow, the uniformity of every SSA value:
+// the widest group of threads within which all threads observe the same value.
+// Operations participate in the analysis by implementing
+// `InferUniformityOpInterface`.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_DATAFLOW_UNIFORMITYANALYSIS_H
+#define MLIR_ANALYSIS_DATAFLOW_UNIFORMITYANALYSIS_H
+
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/Interfaces/InferUniformityOpInterface.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/StringSet.h"
+
+namespace mlir::dataflow {
+
+/// The uniformity lattice element of an SSA value.
+class UniformityLattice : public Lattice<Uniformity> {
+public:
+ using Lattice::Lattice;
+};
+
+/// Uniformity analysis determines, for each SSA value, the widest group of
+/// threads (subgroup, workgroup, cluster, the whole launch) within which every
+/// thread observes the same value. A value that two threads of the same
+/// subgroup may observe differently is divergent.
+///
+/// A value `v` is uniform within a group at a program point `p` if all threads
+/// of the group that execute the dynamic instance of `p` together observe the
+/// same `v`. The analysis only reports uniformity that holds regardless of
+/// which threads are active, so its facts stay true when a value is hoisted
+/// out of divergent control flow.
+///
+/// Transfer functions:
+///
+/// - An operation implementing `InferUniformityOpInterface` describes the
+/// uniformity of the values it defines itself.
+/// - A memory-effect-free operation without regions of a *transparent*
+/// dialect computes a function of its operands, so its results are the join
+/// of the operands (uniform when it has none). The default set of
+/// transparent dialects is the one of the core dialects that have no notion
+/// of a thread.
+/// - Any other operation defines divergent values: an operation that reads
+/// memory, an operation whose regions may capture values from above without
+/// region control flow (a `linalg.generic`), and an operation of a dialect
+/// that may read thread identity without saying so through the interface.
+/// - The results of an operation with region control flow are the join of the
+/// values forwarded to them by region control flow, joined with the
+/// uniformity of the operands that steer that control flow (the condition
+/// of an `scf.if`, the bounds of an `scf.for`, the condition of an
+/// `scf.condition`): threads that take different paths reach the results
+/// with different values, and a loop whose trip count varies across threads
+/// exits with different values. The entry block arguments of the regions
+/// are only the join of the forwarded values, since they are observed by
+/// the threads that entered the region together.
+/// - In a region with unstructured control flow, the arguments of every
+/// non-entry block are joined with the uniformity of the operands that
+/// steer every branch of the region.
+///
+/// The analysis assumes structured reconvergence: all threads of a group that
+/// execute an operation with region control flow together reconverge after
+/// it. This analysis depends on DeadCodeAnalysis and will be a silent no-op if
+/// DeadCodeAnalysis is not loaded in the same solver context.
+class UniformityAnalysis
+ : public SparseForwardDataFlowAnalysis<UniformityLattice> {
+public:
+ /// Creates the analysis. `transparentDialects` lists the dialects whose
+ /// memory-effect-free operations are known to compute a function of their
+ /// operands only; it defaults to `getDefaultTransparentDialects()`.
+ explicit UniformityAnalysis(DataFlowSolver &solver,
+ ArrayRef<StringRef> transparentDialects =
+ getDefaultTransparentDialects());
+
+ /// The dialects treated as transparent by default: `affine`, `arith`,
+ /// `bufferization`, `builtin`, `cf`, `complex`, `func`, `index`, `linalg`,
+ /// `math`, `memref`, `scf`, `tensor`, `ub` and `vector`.
+ static ArrayRef<StringRef> getDefaultTransparentDialects();
+
+ LogicalResult initialize(Operation *top) override;
+
+ LogicalResult visit(ProgramPoint *point) override;
+
+ /// At an entry point, a value is divergent.
+ void setToEntryState(UniformityLattice *lattice) override;
+
+ LogicalResult visitOperation(Operation *op,
+ ArrayRef<const UniformityLattice *> operands,
+ ArrayRef<UniformityLattice *> results) override;
+
+ /// Infers block arguments and results that region control flow does not
+ /// forward a value to (loop induction variables, the body arguments of
+ /// `gpu.launch`) through `InferUniformityOpInterface`.
+ void visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ValueRange nonSuccessorInputs,
+ ArrayRef<UniformityLattice *> nonSuccessorInputLattices) override;
+
+protected:
+ /// Joins the results of a call from its return sites as usual, then adds
+ /// the control dependence of every return site: threads that leave the
+ /// callee through different return sites observe different results.
+ LogicalResult
+ visitCallOperation(CallOpInterface call,
+ ArrayRef<const AbstractSparseLattice *> operandLattices,
+ ArrayRef<AbstractSparseLattice *> resultLattices) override;
+
+ /// Lets a callable implementing `InferUniformityOpInterface` (a GPU kernel)
+ /// describe the uniformity of its arguments.
+ void visitCallableOperation(
+ CallableOpInterface callable,
+ ArrayRef<AbstractSparseLattice *> argLattices) override;
+
+ /// Joins the forwarded values as usual, then adds the control dependence on
+ /// the operands steering the region branch; or defers entirely to
+ /// `InferUniformityOpInterface` when the operation implements it for the
+ /// successor being visited.
+ void
+ visitRegionSuccessors(ProgramPoint *point, RegionBranchOpInterface branch,
+ RegionSuccessor successor,
+ ArrayRef<AbstractSparseLattice *> lattices) override;
+
+private:
+ /// Returns true if `op` belongs to a transparent dialect.
+ bool isTransparent(Operation *op) const;
+
+ /// Runs the interface of `op` at `point` and joins what it reports into the
+ /// lattices of the `candidates` it names. Returns which candidates it set,
+ /// and sets `pending` when an operand of `op` is still uninitialized, in
+ /// which case a candidate the interface did not name must be left alone
+ /// rather than put in the entry state.
+ llvm::SmallBitVector inferThroughInterface(
+ InferUniformityOpInterface op, ProgramPoint *point, ValueRange candidates,
+ ArrayRef<AbstractSparseLattice *> lattices, bool &pending);
+
+ /// Puts the lattices that `set` does not mark in the entry state.
+ void setUnsetToEntryStates(const llvm::SmallBitVector &set,
+ ArrayRef<AbstractSparseLattice *> lattices);
+
+ /// Joins the uniformity of the operands steering `branch` and the
+ /// terminators of its regions into `lattices`, the result lattices of
+ /// `branch`.
+ void joinControlDependence(ProgramPoint *point, Operation *branch,
+ ArrayRef<AbstractSparseLattice *> lattices);
+
+ /// Joins the uniformity of the operands steering every branch in the region
+ /// of `block` into the arguments of `block`, if it is not an entry block.
+ void visitUnstructuredBlockArguments(Block *block);
+
+ llvm::StringSet<> transparentDialects;
+};
+
+/// Collects into `controlOperands` the operands of `op` that steer control
+/// flow rather than being forwarded to a successor: the operands of an
+/// operation implementing `RegionBranchOpInterface` that are not entry
+/// successor operands of any region successor (the condition of `scf.if`, the
+/// bounds of `scf.for`), the operands of a `RegionBranchTerminatorOpInterface`
+/// enclosed in a region branch that are not successor operands of any
+/// successor (the condition of `scf.condition`), and the operands of a
+/// `BranchOpInterface` that are not forwarded to any successor (the condition
+/// of `cf.cond_br`). Collects nothing for any other operation, including a
+/// return-like terminator of a callable.
+void getControlOperands(Operation *op, SmallVectorImpl<Value> &controlOperands);
+
+/// Collects into `controlOperands` the control operands that decide whether a
+/// thread reaches `op`: those of every branch of the regions with unstructured
+/// control flow `op` sits in, and those of every enclosing region branch and
+/// of the terminators of its regions, up to and including `limit` if it is
+/// given, and stopping in any case at the closest enclosing callable or launch
+/// boundary (see `InferUniformityOpInterface::isLaunchBoundary`). The operands
+/// are collected from the innermost to the outermost.
+void getEnclosingControlOperands(Operation *op, Region *limit,
+ SmallVectorImpl<Value> &controlOperands);
+
+/// Returns the uniformity scope of `value` after `solver` has run. A value the
+/// analysis did not reach is reported as divergent.
+UniformityScope getUniformity(DataFlowSolver &solver, Value value);
+
+/// Returns the widest group of threads that execute `op` together: the meet of
+/// the uniformity of the operands steering every enclosing region branch and
+/// every branch of a region with unstructured control flow, up to the closest
+/// enclosing callable or launch boundary (see
+/// `InferUniformityOpInterface::isLaunchBoundary`). A `gpu.barrier` whose
+/// execution uniformity is narrower than its scope may deadlock. If
+/// `narrowingOperand` is given, it receives
+/// the first control operand that narrows the execution to the returned scope,
+/// or null when nothing does.
+UniformityScope getExecutionUniformity(DataFlowSolver &solver, Operation *op,
+ Value *narrowingOperand = nullptr);
+
+} // namespace mlir::dataflow
+
+#endif // MLIR_ANALYSIS_DATAFLOW_UNIFORMITYANALYSIS_H
diff --git a/mlir/lib/Analysis/CMakeLists.txt b/mlir/lib/Analysis/CMakeLists.txt
index 2712ab503663e..3547becaa371d 100644
--- a/mlir/lib/Analysis/CMakeLists.txt
+++ b/mlir/lib/Analysis/CMakeLists.txt
@@ -43,6 +43,7 @@ add_mlir_library(MLIRAnalysis
DataFlow/LivenessAnalysis.cpp
DataFlow/SparseAnalysis.cpp
DataFlow/StridedMetadataRangeAnalysis.cpp
+ DataFlow/UniformityAnalysis.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
@@ -60,6 +61,7 @@ add_mlir_library(MLIRAnalysis
MLIRInferIntRangeInterface
MLIRInferStridedMetadataInterface
MLIRInferTypeOpInterface
+ MLIRInferUniformityOpInterface
MLIRLoopLikeInterface
MLIRPresburger
MLIRSideEffectInterfaces
diff --git a/mlir/lib/Analysis/DataFlow/UniformityAnalysis.cpp b/mlir/lib/Analysis/DataFlow/UniformityAnalysis.cpp
new file mode 100644
index 0000000000000..9657732eaa0f8
--- /dev/null
+++ b/mlir/lib/Analysis/DataFlow/UniformityAnalysis.cpp
@@ -0,0 +1,448 @@
+//===- UniformityAnalysis.cpp - Uniformity analysis -----------------------===//
+//
+// Part of the LLVM Project, 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 "mlir/Analysis/DataFlow/UniformityAnalysis.h"
+
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/CallInterfaces.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "uniformity-analysis"
+
+using namespace mlir;
+using namespace mlir::dataflow;
+
+//===----------------------------------------------------------------------===//
+// Control operands
+//===----------------------------------------------------------------------===//
+
+/// Marks in `forwarded` the positions of the operands in `range`.
+static void markForwarded(OperandRange range, llvm::SmallBitVector &forwarded) {
+ if (range.empty())
+ return;
+ unsigned begin = range.getBeginOperandIndex();
+ for (unsigned i = 0, e = range.size(); i < e; ++i)
+ forwarded.set(begin + i);
+}
+
+void mlir::dataflow::getControlOperands(
+ Operation *op, SmallVectorImpl<Value> &controlOperands) {
+ llvm::SmallBitVector forwarded(op->getNumOperands());
+ if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
+ SmallVector<RegionSuccessor> successors;
+ branch.getSuccessorRegions(RegionBranchPoint::parent(), successors);
+ for (const RegionSuccessor &successor : successors)
+ markForwarded(branch.getEntrySuccessorOperands(successor), forwarded);
+ } else if (auto terminator =
+ dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
+ // The successors of a terminator are defined by the region branch that
+ // encloses it; a return-like terminator of a callable has none.
+ if (!isa<RegionBranchOpInterface>(op->getParentOp()))
+ return;
+ SmallVector<Attribute> operandAttrs(op->getNumOperands(), nullptr);
+ SmallVector<RegionSuccessor> successors;
+ terminator.getSuccessorRegions(operandAttrs, successors);
+ for (const RegionSuccessor &successor : successors)
+ markForwarded(terminator.getSuccessorOperands(successor), forwarded);
+ } else if (auto branch = dyn_cast<BranchOpInterface>(op)) {
+ for (unsigned i = 0, e = op->getNumSuccessors(); i < e; ++i)
+ markForwarded(branch.getSuccessorOperands(i).getForwardedOperands(),
+ forwarded);
+ } else {
+ return;
+ }
+ for (auto [index, operand] : llvm::enumerate(op->getOperands()))
+ if (!forwarded.test(index))
+ controlOperands.push_back(operand);
+}
+
+/// Returns the terminator of `block` if it is one of the branch interfaces
+/// that steer control flow, null otherwise.
+static Operation *getBranchTerminator(Block &block) {
+ if (block.empty())
+ return nullptr;
+ Operation *terminator = &block.back();
+ if (isa<BranchOpInterface, RegionBranchTerminatorOpInterface>(terminator))
+ return terminator;
+ return nullptr;
+}
+
+void mlir::dataflow::getEnclosingControlOperands(
+ Operation *op, Region *limit, SmallVectorImpl<Value> &controlOperands) {
+ for (Operation *current = op;;) {
+ Block *block = current->getBlock();
+ if (!block)
+ return;
+ Region *region = block->getParent();
+ // Unstructured control flow within the region.
+ if (!block->isEntryBlock())
+ for (Block &other : *region)
+ if (Operation *terminator = getBranchTerminator(other))
+ if (isa<BranchOpInterface>(terminator))
+ getControlOperands(terminator, controlOperands);
+ if (region == limit)
+ return;
+
+ Operation *parent = region->getParentOp();
+ if (!parent)
+ return;
+ // The body of a launch is executed by the threads of the launch, whatever
+ // control flow the host wrapped the launch in.
+ if (auto inferrable = dyn_cast<InferUniformityOpInterface>(parent);
+ inferrable && inferrable.isLaunchBoundary())
+ return;
+ if (isa<RegionBranchOpInterface>(parent)) {
+ getControlOperands(parent, controlOperands);
+ for (Region &other : parent->getRegions())
+ for (Block &otherBlock : other)
+ if (Operation *terminator = getBranchTerminator(otherBlock))
+ if (isa<RegionBranchTerminatorOpInterface>(terminator))
+ getControlOperands(terminator, controlOperands);
+ }
+ if (isa<CallableOpInterface>(parent))
+ return;
+ current = parent;
+ }
+}
+
+//===----------------------------------------------------------------------===//
+// UniformityAnalysis
+//===----------------------------------------------------------------------===//
+
+static const StringRef kDefaultTransparentDialects[] = {
+ "affine", "arith", "bufferization", "builtin", "cf",
+ "complex", "func", "index", "linalg", "math",
+ "memref", "scf", "tensor", "ub", "vector"};
+
+ArrayRef<StringRef> UniformityAnalysis::getDefaultTransparentDialects() {
+ return kDefaultTransparentDialects;
+}
+
+UniformityAnalysis::UniformityAnalysis(DataFlowSolver &solver,
+ ArrayRef<StringRef> transparentDialects)
+ : SparseForwardDataFlowAnalysis(solver) {
+ for (StringRef dialect : transparentDialects)
+ this->transparentDialects.insert(dialect);
+}
+
+bool UniformityAnalysis::isTransparent(Operation *op) const {
+ return transparentDialects.contains(op->getName().getDialectNamespace());
+}
+
+void UniformityAnalysis::setToEntryState(UniformityLattice *lattice) {
+ propagateIfChanged(lattice, lattice->join(Uniformity::getDivergent()));
+}
+
+LogicalResult UniformityAnalysis::initialize(Operation *top) {
+ if (failed(SparseForwardDataFlowAnalysis::initialize(top)))
+ return failure();
+ // The base class initializes blocks without going through `visit`, so the
+ // control dependence of unstructured block arguments is added here.
+ top->walk([&](Block *block) { visitUnstructuredBlockArguments(block); });
+ return success();
+}
+
+LogicalResult UniformityAnalysis::visit(ProgramPoint *point) {
+ if (failed(SparseForwardDataFlowAnalysis::visit(point)))
+ return failure();
+ if (point->isBlockStart())
+ visitUnstructuredBlockArguments(point->getBlock());
+ return success();
+}
+
+llvm::SmallBitVector UniformityAnalysis::inferThroughInterface(
+ InferUniformityOpInterface op, ProgramPoint *point, ValueRange candidates,
+ ArrayRef<AbstractSparseLattice *> lattices, bool &pending) {
+ assert(candidates.size() == lattices.size() && "size mismatch");
+ SmallVector<Uniformity> operandUniformity =
+ llvm::map_to_vector(op->getOperands(), [&](Value operand) {
+ return getLatticeElementFor(point, operand)->getValue();
+ });
+ // An implementation is allowed to describe nothing while an operand it needs
+ // is still uninitialized, and several do. That is not the same as saying the
+ // value is divergent: the join is a meet, so a lattice put in the entry state
+ // can never recover, and the operand will settle on a later visit.
+ pending = llvm::any_of(operandUniformity, [](const Uniformity &uniformity) {
+ return uniformity.isUninitialized();
+ });
+ llvm::SmallBitVector set(lattices.size());
+ op.inferUniformity(
+ operandUniformity, [&](Value value, UniformityScope scope) {
+ auto it = llvm::find(candidates, value);
+ if (it == candidates.end())
+ return;
+ unsigned index = std::distance(candidates.begin(), it);
+ set.set(index);
+ auto *lattice = static_cast<UniformityLattice *>(lattices[index]);
+ LDBG() << "Inferred " << scope << " for " << value;
+ propagateIfChanged(lattice, lattice->join(Uniformity(scope)));
+ });
+ return set;
+}
+
+void UniformityAnalysis::setUnsetToEntryStates(
+ const llvm::SmallBitVector &set,
+ ArrayRef<AbstractSparseLattice *> lattices) {
+ for (auto [index, lattice] : llvm::enumerate(lattices))
+ if (!set.test(index))
+ setToEntryState(static_cast<UniformityLattice *>(lattice));
+}
+
+LogicalResult
+UniformityAnalysis::visitOperation(Operation *op,
+ ArrayRef<const UniformityLattice *> operands,
+ ArrayRef<UniformityLattice *> results) {
+ if (auto inferrable = dyn_cast<InferUniformityOpInterface>(op)) {
+ SmallVector<AbstractSparseLattice *> resultLattices(results.begin(),
+ results.end());
+ bool pending = false;
+ llvm::SmallBitVector set =
+ inferThroughInterface(inferrable, getProgramPointAfter(op),
+ op->getResults(), resultLattices, pending);
+ if (!pending)
+ setUnsetToEntryStates(set, resultLattices);
+ return success();
+ }
+
+ // Only a memory-effect-free operation of a transparent dialect that has no
+ // region is known to compute a function of its operands. Anything else may
+ // read memory written by another thread, thread identity, or a value its
+ // region captures from above.
+ if (!isTransparent(op) || !isMemoryEffectFree(op) ||
+ op->getNumRegions() != 0) {
+ setAllToEntryStates(results);
+ return success();
+ }
+
+ Uniformity joined = Uniformity::getUniform();
+ if (!operands.empty()) {
+ joined = Uniformity::join(
+ llvm::map_to_vector(operands, [](const UniformityLattice *lattice) {
+ return lattice->getValue();
+ }));
+ if (joined.isUninitialized())
+ return success();
+ }
+ LDBG() << "Joined operands to " << joined << " for "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ for (UniformityLattice *lattice : results)
+ propagateIfChanged(lattice, lattice->join(joined));
+ return success();
+}
+
+void UniformityAnalysis::visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ValueRange nonSuccessorInputs,
+ ArrayRef<UniformityLattice *> nonSuccessorInputLattices) {
+ if (auto inferrable = dyn_cast<InferUniformityOpInterface>(op)) {
+ ProgramPoint *point =
+ !successor.isRegion()
+ ? getProgramPointAfter(op)
+ : getProgramPointBefore(&successor.getSuccessor()->front());
+ SmallVector<AbstractSparseLattice *> lattices(
+ nonSuccessorInputLattices.begin(), nonSuccessorInputLattices.end());
+ bool pending = false;
+ llvm::SmallBitVector set = inferThroughInterface(
+ inferrable, point, nonSuccessorInputs, lattices, pending);
+ if (pending)
+ return;
+ if (set.any()) {
+ setUnsetToEntryStates(set, lattices);
+ return;
+ }
+ }
+ SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
+ op, successor, nonSuccessorInputs, nonSuccessorInputLattices);
+}
+
+LogicalResult UniformityAnalysis::visitCallOperation(
+ CallOpInterface call,
+ ArrayRef<const AbstractSparseLattice *> operandLattices,
+ ArrayRef<AbstractSparseLattice *> resultLattices) {
+ if (failed(SparseForwardDataFlowAnalysis::visitCallOperation(
+ call, operandLattices, resultLattices)))
+ return failure();
+ if (!getSolverConfig().isInterprocedural())
+ return success();
+ auto callable =
+ dyn_cast_if_present<CallableOpInterface>(call.resolveCallable());
+ Region *body = callable ? callable.getCallableRegion() : nullptr;
+ if (!body)
+ return success();
+ ProgramPoint *point = getProgramPointAfter(call);
+ const auto *predecessors = getOrCreateFor<PredecessorState>(point, point);
+ if (!predecessors->allPredecessorsKnown())
+ return success();
+
+ // The base class joins what every return site yields. Threads that leave
+ // the callee through different return sites observe different results, so
+ // the results also depend on the control flow that selects the return site.
+ SmallVector<Value> controlOperands;
+ for (Operation *returnSite : predecessors->getKnownPredecessors())
+ getEnclosingControlOperands(returnSite, body, controlOperands);
+ Uniformity taint;
+ for (Value operand : controlOperands)
+ taint = Uniformity::join(taint,
+ getLatticeElementFor(point, operand)->getValue());
+ if (taint.isUninitialized())
+ return success();
+ LDBG() << "Return site control dependence " << taint << " for "
+ << OpWithFlags(call, OpPrintingFlags().skipRegions());
+ for (AbstractSparseLattice *lattice : resultLattices) {
+ auto *typed = static_cast<UniformityLattice *>(lattice);
+ propagateIfChanged(typed, typed->join(taint));
+ }
+ return success();
+}
+
+void UniformityAnalysis::visitCallableOperation(
+ CallableOpInterface callable,
+ ArrayRef<AbstractSparseLattice *> argLattices) {
+ auto inferrable =
+ dyn_cast<InferUniformityOpInterface>(callable.getOperation());
+ Region *body = callable.getCallableRegion();
+ if (inferrable && body && !body->empty()) {
+ Block &entry = body->front();
+ bool pending = false;
+ llvm::SmallBitVector set =
+ inferThroughInterface(inferrable, getProgramPointBefore(&entry),
+ entry.getArguments(), argLattices, pending);
+ if (pending)
+ return;
+ if (set.any()) {
+ setUnsetToEntryStates(set, argLattices);
+ return;
+ }
+ }
+ AbstractSparseForwardDataFlowAnalysis::visitCallableOperation(callable,
+ argLattices);
+}
+
+void UniformityAnalysis::visitRegionSuccessors(
+ ProgramPoint *point, RegionBranchOpInterface branch,
+ RegionSuccessor successor, ArrayRef<AbstractSparseLattice *> lattices) {
+ Operation *op = branch.getOperation();
+ if (auto inferrable = dyn_cast<InferUniformityOpInterface>(op)) {
+ // The operation is authoritative for this successor if it sets one of
+ // the inputs that region control flow would otherwise forward a value to.
+ // Setting only a non-successor input (a loop induction variable) leaves
+ // the forwarded ones to the usual join.
+ ValueRange candidates = point->isBlockStart()
+ ? ValueRange(point->getBlock()->getArguments())
+ : ValueRange(op->getResults());
+ bool pending = false;
+ llvm::SmallBitVector set =
+ inferThroughInterface(inferrable, point, candidates, lattices, pending);
+ if (pending)
+ return;
+ ValueRange successorInputs = branch.getSuccessorInputs(successor);
+ bool authoritative = false;
+ for (auto [index, candidate] : llvm::enumerate(candidates))
+ if (set.test(index) && llvm::is_contained(successorInputs, candidate))
+ authoritative = true;
+ if (authoritative) {
+ setUnsetToEntryStates(set, lattices);
+ return;
+ }
+ } else if (!isTransparent(op)) {
+ // A region operation of an unknown dialect may run its regions with any
+ // subset of threads, and its results may come from anywhere. (The typed
+ // overload of setAllToEntryStates hides the untyped one.)
+ AbstractSparseForwardDataFlowAnalysis::setAllToEntryStates(lattices);
+ return;
+ }
+
+ AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors(
+ point, branch, successor, lattices);
+ // Threads that take different paths through the regions reach the results
+ // with different values; the entry block arguments of a region are only
+ // observed by the threads that entered it together.
+ if (!point->isBlockStart())
+ joinControlDependence(point, op, lattices);
+}
+
+void UniformityAnalysis::joinControlDependence(
+ ProgramPoint *point, Operation *branch,
+ ArrayRef<AbstractSparseLattice *> lattices) {
+ SmallVector<Value> controlOperands;
+ getControlOperands(branch, controlOperands);
+ for (Region ®ion : branch->getRegions())
+ for (Block &block : region)
+ if (Operation *terminator = getBranchTerminator(block))
+ getControlOperands(terminator, controlOperands);
+
+ Uniformity taint;
+ for (Value operand : controlOperands)
+ taint = Uniformity::join(taint,
+ getLatticeElementFor(point, operand)->getValue());
+ if (taint.isUninitialized())
+ return;
+ LDBG() << "Control dependence " << taint << " for "
+ << OpWithFlags(branch, OpPrintingFlags().skipRegions());
+ for (AbstractSparseLattice *lattice : lattices) {
+ auto *typed = static_cast<UniformityLattice *>(lattice);
+ propagateIfChanged(typed, typed->join(taint));
+ }
+}
+
+void UniformityAnalysis::visitUnstructuredBlockArguments(Block *block) {
+ if (block->isEntryBlock() || block->getNumArguments() == 0)
+ return;
+ ProgramPoint *point = getProgramPointBefore(block);
+ SmallVector<Value> controlOperands;
+ for (Block &other : *block->getParent())
+ if (Operation *terminator = getBranchTerminator(other))
+ if (isa<BranchOpInterface>(terminator))
+ getControlOperands(terminator, controlOperands);
+
+ Uniformity taint;
+ for (Value operand : controlOperands)
+ taint = Uniformity::join(taint,
+ getLatticeElementFor(point, operand)->getValue());
+ if (taint.isUninitialized())
+ return;
+ for (BlockArgument argument : block->getArguments()) {
+ UniformityLattice *lattice = getLatticeElement(argument);
+ propagateIfChanged(lattice, lattice->join(taint));
+ }
+}
+
+//===----------------------------------------------------------------------===//
+// Queries
+//===----------------------------------------------------------------------===//
+
+UniformityScope mlir::dataflow::getUniformity(DataFlowSolver &solver,
+ Value value) {
+ const auto *lattice = solver.lookupState<UniformityLattice>(value);
+ if (!lattice || lattice->getValue().isUninitialized())
+ return UniformityScope::Divergent;
+ return lattice->getValue().getScope();
+}
+
+UniformityScope
+mlir::dataflow::getExecutionUniformity(DataFlowSolver &solver, Operation *op,
+ Value *narrowingOperand) {
+ SmallVector<Value> controlOperands;
+ getEnclosingControlOperands(op, /*limit=*/nullptr, controlOperands);
+ UniformityScope scope = UniformityScope::Uniform;
+ Value narrowing;
+ for (Value operand : controlOperands) {
+ UniformityScope operandScope = getUniformity(solver, operand);
+ if (operandScope < scope) {
+ scope = operandScope;
+ narrowing = operand;
+ }
+ }
+ if (narrowingOperand)
+ *narrowingOperand = narrowing;
+ return scope;
+}
diff --git a/mlir/test/Analysis/DataFlow/test-uniformity-analysis.mlir b/mlir/test/Analysis/DataFlow/test-uniformity-analysis.mlir
new file mode 100644
index 0000000000000..6a426a3bedc89
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/test-uniformity-analysis.mlir
@@ -0,0 +1,341 @@
+// RUN: mlir-opt -split-input-file -test-uniformity-analysis -verify-diagnostics %s
+
+// A memory-effect-free operation of a transparent dialect is the join of its
+// operands, which is the narrower scope; a constant is uniform.
+func.func @join() {
+ // expected-remark @below {{uniformity of "div": results = [divergent], execution = uniform}}
+ %div = test.with_uniformity {scope = "divergent", tag = "div"} : index
+ %sg = test.with_uniformity {scope = "subgroup"} : index
+ %wg = test.with_uniformity {scope = "workgroup"} : index
+ %cl = test.with_uniformity {scope = "cluster"} : index
+ // expected-remark @below {{uniformity of "uni": results = [uniform], execution = uniform}}
+ %uni = test.with_uniformity {scope = "uniform", tag = "uni"} : index
+ // expected-remark @below {{uniformity of "c": results = [uniform], execution = uniform}}
+ %c = arith.constant {tag = "c"} 4 : index
+ // expected-remark @below {{uniformity of "uni_c": results = [uniform], execution = uniform}}
+ %uni_c = arith.addi %uni, %c {tag = "uni_c"} : index
+ // expected-remark @below {{uniformity of "wg_c": results = [workgroup], execution = uniform}}
+ %wg_c = arith.addi %wg, %c {tag = "wg_c"} : index
+ // expected-remark @below {{uniformity of "wg_sg": results = [subgroup], execution = uniform}}
+ %wg_sg = arith.addi %wg, %sg {tag = "wg_sg"} : index
+ // expected-remark @below {{uniformity of "cl_wg": results = [workgroup], execution = uniform}}
+ %cl_wg = arith.addi %cl, %wg {tag = "cl_wg"} : index
+ // expected-remark @below {{uniformity of "div_uni": results = [divergent], execution = uniform}}
+ %div_uni = arith.addi %div, %uni {tag = "div_uni"} : index
+ // expected-remark @below {{uniformity of "cmp": results = [subgroup], execution = uniform}}
+ %cmp = arith.cmpi ult, %sg, %cl {tag = "cmp"} : index
+ // A select of two uniform values picks the same value within the group
+ // that agrees on the condition, and no wider.
+ // expected-remark @below {{uniformity of "select": results = [subgroup], execution = uniform}}
+ %sel = arith.select %cmp, %uni, %c {tag = "select"} : index
+ return
+}
+
+// -----
+
+// What is not a function of its operands: an operation that reads or allocates
+// memory, an operation whose region captures values from above without region
+// control flow, and an operation of a dialect the analysis was not told is
+// transparent, even a pure one.
+func.func @opaque() {
+ %c0 = arith.constant 0 : index
+ %uni = test.with_uniformity {scope = "uniform"} : index
+ // expected-remark @below {{uniformity of "alloca": results = [divergent], execution = uniform}}
+ %m = memref.alloca() {tag = "alloca"} : memref<4xindex>
+ // expected-remark @below {{uniformity of "load": results = [divergent], execution = uniform}}
+ %l = memref.load %m[%c0] {tag = "load"} : memref<4xindex>
+ // expected-remark @below {{uniformity of "generate": results = [divergent], execution = uniform}}
+ %t = tensor.generate {
+ ^bb0(%i: index):
+ tensor.yield %uni : index
+ } {tag = "generate"} : tensor<4xindex>
+ // expected-remark @below {{uniformity of "not_transparent": results = [divergent], execution = uniform}}
+ %r = test.increment {tag = "not_transparent"} %uni : index
+ return
+}
+
+// -----
+
+// Structured control flow. The results of a region branch are the join of the
+// values forwarded to them, tainted by the operands steering the branch; the
+// entry block arguments of a region are only the join of the forwarded values.
+// An operation executes as widely as the narrowest control operand around it.
+func.func @structured() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %uni_flag = test.with_uniformity {scope = "uniform"} : i1
+ %wg_flag = test.with_uniformity {scope = "workgroup"} : i1
+ %div_flag = test.with_uniformity {scope = "divergent"} : i1
+ %n = test.with_uniformity {scope = "uniform"} : index
+ %tid = test.with_uniformity {scope = "divergent"} : index
+ // expected-remark @below {{uniformity of "if_div": results = [divergent], execution = uniform}}
+ %r0 = scf.if %div_flag -> index {
+ // expected-remark @below {{uniformity of "in_then": results = [uniform], execution = divergent}}
+ %a = arith.addi %c1, %c1 {tag = "in_then"} : index
+ scf.yield %a : index
+ } else {
+ scf.yield %c0 : index
+ } {tag = "if_div"}
+ // expected-remark @below {{uniformity of "if_uni": results = [uniform], execution = uniform}}
+ %r1 = scf.if %uni_flag -> index {
+ scf.yield %c1 : index
+ } else {
+ scf.yield %c0 : index
+ } {tag = "if_uni"}
+ // expected-remark @below {{uniformity of "if_wg": results = [workgroup], execution = uniform}}
+ %r2 = scf.if %wg_flag -> index {
+ // expected-remark @below {{uniformity of "in_wg_then": results = [uniform], execution = workgroup}}
+ %a = arith.addi %c1, %c1 {tag = "in_wg_then"} : index
+ scf.yield %a : index
+ } else {
+ scf.yield %c0 : index
+ } {tag = "if_wg"}
+ // A loop with uniform bounds has uniform results, and its iteration argument
+ // is the join of the initial value and the yielded value.
+ // expected-remark @below {{uniformity of "for_uni": results = [uniform], execution = uniform}}
+ %r3 = scf.for %i = %c0 to %n step %c1 iter_args(%acc = %c0) -> index {
+ // expected-remark @below {{uniformity of "acc_uni": results = [uniform], execution = uniform}}
+ %next = arith.addi %acc, %c1 {tag = "acc_uni"} : index
+ scf.yield %next : index
+ } {tag = "for_uni"}
+ // A loop whose trip count differs between threads exits with different
+ // values and its body executes divergently, but the threads still iterating
+ // observe the same iteration argument.
+ // expected-remark @below {{uniformity of "for_div_ub": results = [divergent], execution = uniform}}
+ %r4 = scf.for %j = %c0 to %tid step %c1 iter_args(%acc = %c0) -> index {
+ // expected-remark @below {{uniformity of "acc_div_ub": results = [uniform], execution = divergent}}
+ %next = arith.addi %acc, %c1 {tag = "acc_div_ub"} : index
+ scf.yield %next : index
+ } {tag = "for_div_ub"}
+ // Nested branches: the narrowest control operand wins.
+ scf.if %wg_flag {
+ // expected-remark @below {{uniformity of "in_wg_if": results = [uniform], execution = workgroup}}
+ %a = arith.constant {tag = "in_wg_if"} 1 : index
+ scf.if %div_flag {
+ // expected-remark @below {{uniformity of "nested": results = [uniform], execution = divergent}}
+ %b = arith.constant {tag = "nested"} 1 : index
+ }
+ scf.if %uni_flag {
+ // expected-remark @below {{uniformity of "nested_uni": results = [uniform], execution = workgroup}}
+ %d = arith.constant {tag = "nested_uni"} 1 : index
+ }
+ }
+ return
+}
+
+// -----
+
+// scf.while: the condition steers the results and the execution of the body.
+func.func @while() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %n = test.with_uniformity {scope = "uniform"} : index
+ %tid = test.with_uniformity {scope = "divergent"} : index
+ // expected-remark @below {{uniformity of "while_div": results = [divergent], execution = uniform}}
+ %r = scf.while (%i = %c0) : (index) -> index {
+ %cond = arith.cmpi ult, %i, %tid : index
+ scf.condition(%cond) %i : index
+ } do {
+ ^bb0(%i: index):
+ // expected-remark @below {{uniformity of "in_while": results = [uniform], execution = divergent}}
+ %next = arith.addi %i, %c1 {tag = "in_while"} : index
+ scf.yield %next : index
+ } attributes {tag = "while_div"}
+ // expected-remark @below {{uniformity of "while_uni": results = [uniform], execution = uniform}}
+ %s = scf.while (%i = %c0) : (index) -> index {
+ %cond = arith.cmpi ult, %i, %n : index
+ scf.condition(%cond) %i : index
+ } do {
+ ^bb0(%i: index):
+ // expected-remark @below {{uniformity of "in_while_uni": results = [uniform], execution = uniform}}
+ %next = arith.addi %i, %c1 {tag = "in_while_uni"} : index
+ scf.yield %next : index
+ } attributes {tag = "while_uni"}
+ return
+}
+
+// -----
+
+// Unstructured control flow: the arguments of a non-entry block are tainted by
+// every branch of the region, and so is their execution.
+func.func @cfg() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cond = test.with_uniformity {scope = "divergent"} : i1
+ cf.cond_br %cond, ^bb1(%c0 : index), ^bb1(%c1 : index)
+^bb1(%x: index):
+ // expected-remark @below {{uniformity of "phi": results = [divergent], execution = divergent}}
+ %y = arith.addi %x, %c1 {tag = "phi"} : index
+ return
+}
+
+func.func @cfg_uniform() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cond = test.with_uniformity {scope = "uniform"} : i1
+ cf.cond_br %cond, ^bb1(%c0 : index), ^bb1(%c1 : index)
+^bb1(%x: index):
+ // expected-remark @below {{uniformity of "phi_uni": results = [uniform], execution = uniform}}
+ %y = arith.addi %x, %c1 {tag = "phi_uni"} : index
+ return
+}
+
+// -----
+
+// A launch boundary: its body is executed by the threads of the launch,
+// whatever control flow the host wrapped the launch in, and the operation
+// describes the arguments of its body.
+func.func @launch() {
+ %c1 = arith.constant 1 : index
+ %go = test.with_uniformity {scope = "divergent"} : i1
+ scf.if %go {
+ test.uniformity_launch ["workgroup", "divergent"] {
+ ^bb0(%block: index, %thread: index):
+ // expected-remark @below {{uniformity of "launch_block": results = [workgroup], execution = uniform}}
+ %a = arith.addi %block, %c1 {tag = "launch_block"} : index
+ // expected-remark @below {{uniformity of "launch_thread": results = [divergent], execution = uniform}}
+ %b = arith.addi %thread, %c1 {tag = "launch_thread"} : index
+ }
+ }
+ return
+}
+
+// -----
+
+// A value an operation could not describe yet, because an operand had not
+// been visited, must not be pinned to divergent: the loop result below is only
+// known once the regions of the while have settled, and the operation using
+// it is visited first.
+func.func @late_operand() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %n = test.with_uniformity {scope = "uniform"} : index
+ %r = scf.while (%i = %c0) : (index) -> index {
+ %c = arith.cmpi ult, %i, %n : index
+ scf.condition(%c) %i : index
+ } do {
+ ^bb0(%i: index):
+ %next = arith.addi %i, %c1 : index
+ scf.yield %next : index
+ }
+ // expected-remark @below {{uniformity of "late": results = [uniform, divergent], execution = uniform}}
+ %same, %other = test.uniformity_of %r {tag = "late"} : index
+ return
+}
+
+// -----
+
+// A callable the analysis is not told about receives divergent arguments: it
+// joins every call site, and nothing says a function is only called from
+// uniform control flow. Execution uniformity stops at the callable, so the
+// body of a callee is reported as executed by every thread whatever the call
+// site.
+func.func @callee(%x: index) -> index {
+ // expected-remark @below {{uniformity of "callee_arg": results = [divergent], execution = uniform}}
+ %a = arith.addi %x, %x {tag = "callee_arg"} : index
+ return %a : index
+}
+
+func.func @caller() {
+ %n = test.with_uniformity {scope = "uniform"} : index
+ %go = test.with_uniformity {scope = "divergent"} : i1
+ // expected-remark @below {{uniformity of "call": results = [divergent], execution = uniform}}
+ %r = func.call @callee(%n) {tag = "call"} : (index) -> index
+ scf.if %go {
+ %s = func.call @callee(%n) : (index) -> index
+ }
+ return
+}
+
+// -----
+
+// A region branch that describes its successor inputs itself, like
+// gpu.warp_execute_on_lane_0. Setting one input of a successor makes the
+// operation authoritative for that successor: the forwarded values are not
+// joined into its inputs, and the inputs it leaves undescribed are put in the
+// entry state. A successor the operation says nothing about is the usual join,
+// tainted by the control operand for the results only.
+func.func @authoritative() {
+ %div = test.with_uniformity {scope = "divergent"} : index
+ %wg = test.with_uniformity {scope = "workgroup"} : index
+ %uni = test.with_uniformity {scope = "uniform"} : index
+ %uni_flag = test.with_uniformity {scope = "uniform"} : i1
+ %wg_flag = test.with_uniformity {scope = "workgroup"} : i1
+ // The result is declared subgroup although the region yields a divergent
+ // value; the body argument, which the operation does not describe, is the
+ // join of the divergent value forwarded to it.
+ // expected-remark @below {{uniformity of "result_set": results = [subgroup], execution = uniform}}
+ %r0 = test.uniformity_region %uni_flag(%div) results ["subgroup"] : (index) -> index {
+ ^bb0(%x: index):
+ // expected-remark @below {{uniformity of "result_set_arg": results = [divergent], execution = uniform}}
+ %y = arith.addi %x, %x {tag = "result_set_arg"} : index
+ test.region_if_yield %y : index
+ } {tag = "result_set"}
+ // The body argument is declared uniform although a divergent value is
+ // forwarded to it; the result, which the operation does not describe, is
+ // the join of the yielded value.
+ // expected-remark @below {{uniformity of "arg_set": results = [workgroup], execution = uniform}}
+ %r1 = test.uniformity_region %uni_flag(%div) args ["uniform"] : (index) -> index {
+ ^bb0(%x: index):
+ // expected-remark @below {{uniformity of "arg_set_arg": results = [uniform], execution = uniform}}
+ %y = arith.addi %x, %x {tag = "arg_set_arg"} : index
+ test.region_if_yield %wg : index
+ } {tag = "arg_set"}
+ // Nothing is described: the result is the join of the uniform yielded value
+ // and of the workgroup control operand, the body argument is only the join
+ // of the forwarded value, and the body executes within the workgroup.
+ // expected-remark @below {{uniformity of "unset": results = [workgroup], execution = uniform}}
+ %r2 = test.uniformity_region %wg_flag(%uni) : (index) -> index {
+ ^bb0(%x: index):
+ // expected-remark @below {{uniformity of "unset_arg": results = [uniform], execution = workgroup}}
+ %y = arith.addi %x, %x {tag = "unset_arg"} : index
+ test.region_if_yield %y : index
+ } {tag = "unset"}
+ // Setting one result makes the operation authoritative for both: the one it
+ // leaves undescribed is put in the entry state, not joined with the uniform
+ // value yielded to it.
+ // expected-remark @below {{uniformity of "partial": results = [subgroup, divergent], execution = uniform}}
+ %r3:2 = test.uniformity_region %uni_flag(%uni) results ["subgroup", ""] : (index) -> (index, index) {
+ ^bb0(%x: index):
+ test.region_if_yield %x, %x : index, index
+ } {tag = "partial"}
+ return
+}
+
+// -----
+
+// The results of a call are the join of what every return site yields, and
+// threads that leave the callee through different return sites observe
+// different results: the control flow that selects the return site taints
+// them, even when every return site yields a uniform value. The callees are
+// private so that every call site is known and their arguments are the join
+// of what the callers pass.
+func.func private @two_returns(%c: i1) -> index {
+ %one = arith.constant 1 : index
+ %two = arith.constant 2 : index
+ cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+ return %one : index
+^bb2:
+ return %two : index
+}
+
+func.func private @two_returns_uniform(%c: i1) -> index {
+ %one = arith.constant 1 : index
+ %two = arith.constant 2 : index
+ cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+ return %one : index
+^bb2:
+ return %two : index
+}
+
+func.func @call_return_sites() {
+ %div = test.with_uniformity {scope = "divergent"} : i1
+ %uni = test.with_uniformity {scope = "uniform"} : i1
+ // expected-remark @below {{uniformity of "call_div_sites": results = [divergent], execution = uniform}}
+ %r0 = func.call @two_returns(%div) {tag = "call_div_sites"} : (i1) -> index
+ // expected-remark @below {{uniformity of "call_uni_sites": results = [uniform], execution = uniform}}
+ %r1 = func.call @two_returns_uniform(%uni) {tag = "call_uni_sites"} : (i1) -> index
+ return
+}
diff --git a/mlir/test/lib/Analysis/CMakeLists.txt b/mlir/test/lib/Analysis/CMakeLists.txt
index d86af5017f24b..1b69f7aa9fd31 100644
--- a/mlir/test/lib/Analysis/CMakeLists.txt
+++ b/mlir/test/lib/Analysis/CMakeLists.txt
@@ -19,6 +19,7 @@ add_mlir_library(MLIRTestAnalysis
DataFlow/TestLivenessAnalysis.cpp
DataFlow/TestSparseBackwardDataFlowAnalysis.cpp
DataFlow/TestStridedMetadataRangeAnalysis.cpp
+ DataFlow/TestUniformityAnalysis.cpp
EXCLUDE_FROM_LIBMLIR
diff --git a/mlir/test/lib/Analysis/DataFlow/TestUniformityAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestUniformityAnalysis.cpp
new file mode 100644
index 0000000000000..c6dbb9bc99f92
--- /dev/null
+++ b/mlir/test/lib/Analysis/DataFlow/TestUniformityAnalysis.cpp
@@ -0,0 +1,66 @@
+//===- TestUniformityAnalysis.cpp - Test uniformity analysis --------------===//
+//
+// Part of the LLVM Project, 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 "mlir/Analysis/DataFlow/UniformityAnalysis.h"
+#include "mlir/Analysis/DataFlow/Utils.h"
+#include "mlir/Analysis/DataFlowFramework.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+using namespace mlir::dataflow;
+
+namespace {
+/// Reports the uniformity of the results of every operation carrying a `tag`
+/// string attribute, and the uniformity of its execution, as a remark.
+struct TestUniformityAnalysisPass
+ : public PassWrapper<TestUniformityAnalysisPass, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestUniformityAnalysisPass)
+
+ StringRef getArgument() const override { return "test-uniformity-analysis"; }
+ StringRef getDescription() const override {
+ return "Test uniformity analysis by reporting the uniformity of the "
+ "results and of the execution of every operation with a `tag` "
+ "attribute";
+ }
+
+ void runOnOperation() override {
+ Operation *rootOp = getOperation();
+
+ DataFlowSolver solver;
+ loadBaselineAnalyses(solver);
+ solver.load<UniformityAnalysis>();
+ if (failed(solver.initializeAndRun(rootOp)))
+ return signalPassFailure();
+
+ rootOp->walk([&](Operation *op) {
+ auto tag = op->getAttrOfType<StringAttr>("tag");
+ if (!tag)
+ return;
+ std::string message;
+ llvm::raw_string_ostream os(message);
+ os << "uniformity of \"" << tag.getValue() << "\": results = [";
+ llvm::interleaveComma(op->getResults(), os, [&](Value result) {
+ os << getUniformity(solver, result);
+ });
+ os << "], execution = " << getExecutionUniformity(solver, op);
+ op->emitRemark() << message;
+ });
+ }
+};
+} // end anonymous namespace
+
+namespace mlir::test {
+void registerTestUniformityAnalysisPass() {
+ PassRegistration<TestUniformityAnalysisPass>();
+}
+} // end namespace mlir::test
diff --git a/mlir/test/lib/Dialect/Test/CMakeLists.txt b/mlir/test/lib/Dialect/Test/CMakeLists.txt
index 4c4721193322b..ee3413b21ca0d 100644
--- a/mlir/test/lib/Dialect/Test/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/Test/CMakeLists.txt
@@ -85,6 +85,7 @@ mlir_target_link_libraries(MLIRTestDialect PUBLIC
MLIRIR
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
+ MLIRInferUniformityOpInterface
MLIRLinalgDialect
MLIRLinalgTransforms
MLIRPtrDialect
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 06871fd0ea835..d37b1fb4a45f7 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -12,6 +12,7 @@
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/FunctionImplementation.h"
+#include "mlir/Interfaces/InferUniformityOpInterface.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
#include "llvm/ADT/SmallVectorExtras.h"
@@ -948,6 +949,127 @@ void TestWithoutBoundsOp::inferResultRangesFromOptional(
setResultRanges(getResult(), IntegerValueRange{});
}
+//===----------------------------------------------------------------------===//
+// TestWithUniformityOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult TestWithUniformityOp::verify() {
+ if (!symbolizeUniformityScope(getScope()))
+ return emitOpError("unknown uniformity scope \"") << getScope() << "\"";
+ return success();
+}
+
+void TestWithUniformityOp::inferUniformity(ArrayRef<Uniformity>,
+ SetUniformityFn setUniformity) {
+ setUniformity(getResult(), *symbolizeUniformityScope(getScope()));
+}
+
+//===----------------------------------------------------------------------===//
+// TestUniformityOfOp
+//===----------------------------------------------------------------------===//
+
+void TestUniformityOfOp::inferUniformity(ArrayRef<Uniformity> operandUniformity,
+ SetUniformityFn setUniformity) {
+ if (!operandUniformity[0].isUninitialized())
+ setUniformity(getSame(), operandUniformity[0].getScope());
+ setUniformity(getOther(), UniformityScope::Divergent);
+}
+
+//===----------------------------------------------------------------------===//
+// TestUniformityLaunchOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult TestUniformityLaunchOp::verify() {
+ if (getScopes().size() != getBody().getNumArguments())
+ return emitOpError("expected one scope per body argument");
+ for (Attribute scope : getScopes())
+ if (!symbolizeUniformityScope(cast<StringAttr>(scope).getValue()))
+ return emitOpError("unknown uniformity scope ") << scope;
+ return success();
+}
+
+void TestUniformityLaunchOp::inferUniformity(ArrayRef<Uniformity>,
+ SetUniformityFn setUniformity) {
+ for (auto [argument, scope] :
+ llvm::zip_equal(getBody().getArguments(), getScopes()))
+ setUniformity(argument, *symbolizeUniformityScope(
+ cast<StringAttr>(scope).getValue()));
+}
+
+bool TestUniformityLaunchOp::isLaunchBoundary() { return true; }
+
+//===----------------------------------------------------------------------===//
+// TestUniformityRegionOp
+//===----------------------------------------------------------------------===//
+
+/// Checks that `scopes`, if present, has one entry per value, each a known
+/// scope or the empty string.
+static LogicalResult verifyUniformityScopes(Operation *op,
+ std::optional<ArrayAttr> scopes,
+ unsigned numValues,
+ StringRef what) {
+ if (!scopes)
+ return success();
+ if (scopes->size() != numValues)
+ return op->emitOpError("expected one scope per ") << what;
+ for (Attribute scope : *scopes) {
+ StringRef name = cast<StringAttr>(scope).getValue();
+ if (!name.empty() && !symbolizeUniformityScope(name))
+ return op->emitOpError("unknown uniformity scope ") << scope;
+ }
+ return success();
+}
+
+LogicalResult TestUniformityRegionOp::verify() {
+ if (failed(verifyUniformityScopes(getOperation(), getArgScopes(),
+ getBody().getNumArguments(),
+ "body argument")))
+ return failure();
+ return verifyUniformityScopes(getOperation(), getResultScopes(),
+ getNumResults(), "result");
+}
+
+void TestUniformityRegionOp::getSuccessorRegions(
+ RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {
+ // The body is always entered, and returns to the parent.
+ if (point.isParent())
+ regions.emplace_back(&getBody());
+ else
+ regions.push_back(RegionSuccessor(getOperation()));
+}
+
+OperandRange
+TestUniformityRegionOp::getEntrySuccessorOperands(RegionSuccessor) {
+ return getInits();
+}
+
+ValueRange
+TestUniformityRegionOp::getSuccessorInputs(RegionSuccessor successor) {
+ if (successor.isOperation())
+ return getResults();
+ return getBody().getArguments();
+}
+
+/// Declares the scope listed in `scopes` for each of `values`, skipping the
+/// values whose entry is the empty string.
+static void setUniformityScopes(std::optional<ArrayAttr> scopes,
+ ValueRange values,
+ SetUniformityFn setUniformity) {
+ if (!scopes)
+ return;
+ for (auto [value, scope] : llvm::zip_equal(values, *scopes)) {
+ StringRef name = cast<StringAttr>(scope).getValue();
+ if (!name.empty())
+ setUniformity(value, *symbolizeUniformityScope(name));
+ }
+}
+
+void TestUniformityRegionOp::inferUniformity(ArrayRef<Uniformity>,
+ SetUniformityFn setUniformity) {
+ setUniformityScopes(getArgScopes(), getBody().getArguments(), setUniformity);
+ setUniformityScopes(getResultScopes(), getResults(), setUniformity);
+}
+
//===----------------------------------------------------------------------===//
// TestWithBoundsRegionOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.h b/mlir/test/lib/Dialect/Test/TestOps.h
index 3831301721990..0e1797e47a3e7 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.h
+++ b/mlir/test/lib/Dialect/Test/TestOps.h
@@ -36,6 +36,7 @@
#include "mlir/Interfaces/DerivedAttributeOpInterface.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/InferUniformityOpInterface.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index f966025dfded3..7f0435409f259 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -27,6 +27,7 @@ include "mlir/Interfaces/DataLayoutInterfaces.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/InferUniformityOpInterface.td"
include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/MemorySlotInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -3565,6 +3566,108 @@ def TestReflectBoundsOp : TEST_Op<"reflect_bounds",
let assemblyFormat = "attr-dict $value `:` type($result)";
}
+//===----------------------------------------------------------------------===//
+// Test InferUniformityOpInterface
+//===----------------------------------------------------------------------===//
+
+def TestWithUniformityOp : TEST_Op<"with_uniformity",
+ [DeclareOpInterfaceMethods<InferUniformityOpInterface>, NoMemoryEffect]> {
+ let description = [{
+ Defines a value of the given uniformity scope, for uniformity analysis
+ tests.
+
+ Example:
+
+ ```mlir
+ %0 = test.with_uniformity {scope = "workgroup"} : index
+ ```
+ }];
+ let arguments = (ins StrAttr:$scope);
+ let results = (outs AnyType:$result);
+ let assemblyFormat = "attr-dict `:` type($result)";
+ let hasVerifier = 1;
+}
+
+def TestUniformityOfOp : TEST_Op<"uniformity_of",
+ [DeclareOpInterfaceMethods<InferUniformityOpInterface>, NoMemoryEffect,
+ AllTypesMatch<["value", "same", "other"]>]> {
+ let description = [{
+ Defines two values from one: `same` has the uniformity of the operand and
+ `other` is divergent, like the value and the validity flag of a shuffle.
+ Says nothing about `same` while the operand is uninitialized.
+
+ Example:
+
+ ```mlir
+ %same, %other = test.uniformity_of %0 : index
+ ```
+ }];
+ let arguments = (ins AnyType:$value);
+ let results = (outs AnyType:$same, AnyType:$other);
+ let assemblyFormat = "$value attr-dict `:` type($value)";
+}
+
+def TestUniformityLaunchOp : TEST_Op<"uniformity_launch",
+ [DeclareOpInterfaceMethods<InferUniformityOpInterface,
+ ["isLaunchBoundary"]>,
+ SingleBlock, NoTerminator, RecursiveMemoryEffects]> {
+ let description = [{
+ A launch boundary: its region is executed by the threads of a launch
+ whatever control flow encloses the operation, and the entry block arguments
+ of the region have the uniformity scopes listed.
+
+ Example:
+
+ ```mlir
+ test.uniformity_launch ["workgroup", "divergent"] {
+ ^bb0(%block: index, %thread: index):
+ ...
+ }
+ ```
+ }];
+ let arguments = (ins StrArrayAttr:$scopes);
+ let regions = (region SizedRegion<1>:$body);
+ let assemblyFormat = "$scopes attr-dict-with-keyword $body";
+ let hasVerifier = 1;
+}
+
+def TestUniformityRegionOp : TEST_Op<"uniformity_region",
+ [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands",
+ "getSuccessorInputs"]>,
+ DeclareOpInterfaceMethods<InferUniformityOpInterface>, SingleBlock,
+ RecursiveMemoryEffects]> {
+ let description = [{
+ A region branch that describes its successor inputs itself, like
+ `gpu.warp_execute_on_lane_0`: `cond` steers the branch, `inits` are
+ forwarded to the arguments of the body and the operands of its
+ `test.region_if_yield` terminator to the results. `arg_scopes` (`result_scopes`), when present, lists in
+ order the uniformity scope the operation declares for every body argument
+ (result); an empty string leaves that value undescribed. When absent, the
+ operation says nothing about that successor.
+
+ Example:
+
+ ```mlir
+ %r = test.uniformity_region %cond(%a) args ["uniform"] results ["subgroup"]
+ : (index) -> index {
+ ^bb0(%x: index):
+ test.region_if_yield %x : index
+ }
+ ```
+ }];
+ let arguments = (ins I1:$cond, Variadic<AnyType>:$inits,
+ OptionalAttr<StrArrayAttr>:$arg_scopes,
+ OptionalAttr<StrArrayAttr>:$result_scopes);
+ let results = (outs Variadic<AnyType>:$results);
+ let regions = (region SizedRegion<1>:$body);
+ let assemblyFormat = [{
+ $cond `(` $inits `)` (`args` $arg_scopes^)? (`results` $result_scopes^)?
+ `:` functional-type($inits, $results) $body attr-dict
+ }];
+ let hasVerifier = 1;
+}
+
//===----------------------------------------------------------------------===//
// Test ConditionallySpeculatable
//===----------------------------------------------------------------------===//
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 47dffb1b2d28b..8854cf9b35da9 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -158,6 +158,7 @@ void registerTestTensorTransforms();
void registerTestTopologicalSortAnalysisPass();
void registerTestTransformDialectEraseSchedulePass();
void registerTestPassStateExtensionCommunication();
+void registerTestUniformityAnalysisPass();
void registerTestVectorLowerings();
void registerTestVectorReductionToSPIRVDotProd();
void registerTestVulkanRunnerPipeline();
@@ -307,6 +308,7 @@ static void registerTestPasses() {
mlir::test::registerTestTopologicalSortAnalysisPass();
mlir::test::registerTestTransformDialectEraseSchedulePass();
mlir::test::registerTestPassStateExtensionCommunication();
+ mlir::test::registerTestUniformityAnalysisPass();
mlir::test::registerTestVectorLowerings();
mlir::test::registerTestVectorReductionToSPIRVDotProd();
mlir::test::registerTestVulkanRunnerPipeline();
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 2efe8697a6751..a760379f2eb41 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -9147,6 +9147,7 @@ cc_library(
":InferIntDivisibilityOpInterface",
":InferIntRangeInterface",
":InferStridedMetadataInterface",
+ ":InferUniformityOpInterface",
":LoopLikeInterface",
":Pass",
":SideEffectInterfaces",
@@ -13286,6 +13287,7 @@ cc_library(
":InferIntRangeInterface",
":InferStridedMetadataInterface",
":InferTypeOpInterface",
+ ":InferUniformityOpInterface",
":InliningUtils",
":Pass",
":SideEffectInterfaces",
diff --git a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
index 22e38c5465c81..c31d0778a1a64 100644
--- a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
@@ -138,6 +138,7 @@ td_library(
"//mlir:FunctionInterfacesTdFiles",
"//mlir:InferIntRangeInterfaceTdFiles",
"//mlir:InferTypeOpInterfaceTdFiles",
+ "//mlir:InferUniformityOpInterfaceTdFiles",
"//mlir:LinalgStructuredOpsTdFiles",
"//mlir:MemorySlotInterfacesTdFiles",
"//mlir:NVVMRequiresSMTraitsTdFiles",
@@ -424,6 +425,7 @@ cc_library(
"//mlir:IR",
"//mlir:InferIntRangeInterface",
"//mlir:InferTypeOpInterface",
+ "//mlir:InferUniformityOpInterface",
"//mlir:InliningUtils",
"//mlir:LLVMDialect",
"//mlir:LLVMIRToLLVMTranslation",
More information about the llvm-commits
mailing list