[Mlir-commits] [mlir] dc520a5 - [mlir][GPU] Add ValueBoundsOphinterface to gpu.subgroup_broadcast (#183848)

[llvmlistbot at llvm.org]

Author: Krzysztof Drewniak
Date: 2026-02-27T15:46:56-08:00
New Revision: dc520a5f493aac270315304be10b9b8bc5119dfe

URL: https://github.com/llvm/llvm-project/commit/dc520a5f493aac270315304be10b9b8bc5119dfe
DIFF: https://github.com/llvm/llvm-project/commit/dc520a5f493aac270315304be10b9b8bc5119dfe.diff

LOG: [mlir][GPU] Add  ValueBoundsOphinterface to gpu.subgroup_broadcast (#183848)

This commit adds an ValueBoundsOpInterface to gpu.subgroup_broadcast,
matching its integer range interface implementation, so that affine
analysis can peek through subgroup broadcast ops.

Added: 
    

Modified: 
    mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
    mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
    mlir/test/Dialect/GPU/value-bounds-op-interface-impl.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index a66a83b7e3ca1..d347cda64a9f0 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -283,10 +283,11 @@ void GPUDialect::initialize() {
   addInterfaces<GPUInlinerInterface>();
   declarePromisedInterface<bufferization::BufferDeallocationOpInterface,
                            TerminatorOp>();
-  declarePromisedInterfaces<
-      ValueBoundsOpInterface, ClusterDimOp, ClusterDimBlocksOp, ClusterIdOp,
-      ClusterBlockIdOp, BlockDimOp, BlockIdOp, GridDimOp, ThreadIdOp, LaneIdOp,
-      SubgroupIdOp, GlobalIdOp, NumSubgroupsOp, SubgroupSizeOp, LaunchOp>();
+  declarePromisedInterfaces<ValueBoundsOpInterface, ClusterDimOp,
+                            ClusterDimBlocksOp, ClusterIdOp, ClusterBlockIdOp,
+                            BlockDimOp, BlockIdOp, GridDimOp, ThreadIdOp,
+                            LaneIdOp, SubgroupIdOp, GlobalIdOp, NumSubgroupsOp,
+                            SubgroupSizeOp, LaunchOp, SubgroupBroadcastOp>();
 }
 
 static std::string getSparseHandleKeyword(SparseHandleKind kind) {

diff  --git a/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
index 3bb7082daa5a0..c978aefe3f655 100644
--- a/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -40,6 +40,28 @@ struct GpuIdOpInterface
   }
 };
 
+/// Implement ValueBoundsOpInterface on subgroup broadcast operations to
+/// indicate that such a broadcast does not modify the ranges of the values in
+/// question. Handles shaped types just in case one wants to broadcast a memref
+/// descriptor.
+struct SubgroupBroadcastOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<SubgroupBroadcastOpInterface,
+                                                   SubgroupBroadcastOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto broadcastOp = cast<SubgroupBroadcastOp>(op);
+    assert(value == broadcastOp.getResult() && "invalid value");
+    cstr.bound(value) == cstr.getExpr(broadcastOp.getSrc());
+  }
+
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto broadcastOp = cast<SubgroupBroadcastOp>(op);
+    assert(value == broadcastOp.getResult() && "invalid value");
+    cstr.bound(value)[dim] == cstr.getExpr(broadcastOp.getSrc(), dim);
+  }
+};
+
 struct GpuLaunchOpInterface
     : public ValueBoundsOpInterface::ExternalModel<GpuLaunchOpInterface,
                                                    LaunchOp> {
@@ -110,5 +132,6 @@ void mlir::gpu::registerValueBoundsOpInterfaceExternalModels(
 #undef REGISTER
 
     LaunchOp::attachInterface<GpuLaunchOpInterface>(*ctx);
+    SubgroupBroadcastOp::attachInterface<SubgroupBroadcastOpInterface>(*ctx);
   });
 }

diff  --git a/mlir/test/Dialect/GPU/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/GPU/value-bounds-op-interface-impl.mlir
index 6facf1e22aab9..d6a2e17e19704 100644
--- a/mlir/test/Dialect/GPU/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/GPU/value-bounds-op-interface-impl.mlir
@@ -157,3 +157,41 @@ module attributes {gpu.container_module} {
     }
   }
 }
+
+// -----
+
+// CHECK-LABEL: func @subgroup_broadcast
+module attributes {gpu.container_module} {
+  gpu.module @gpu_module {
+    gpu.func @subgroup_broadcast(%arg0 : index) kernel {
+      %lane = arith.constant 1 : i32
+
+      %bcast = gpu.subgroup_broadcast %arg0, specific_lane %lane : index
+      // expected-remark @below{{true}}
+      "test.compare"(%bcast, %arg0) {cmp = "EQ"} : (index, index) -> ()
+
+      %bcast2 = gpu.subgroup_broadcast %arg0, first_active_lane : index
+      // expected-remark @below{{true}}
+      "test.compare"(%bcast2, %arg0) {cmp = "EQ"} : (index, index) -> ()
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+// CHECK-LABEL: func @subgroup_broadcast_shaped
+module attributes {gpu.container_module} {
+  gpu.module @gpu_module {
+    gpu.func @subgroup_broadcast_shaped(%arg0 : memref<?xf32>) kernel {
+      %c0 = arith.constant 0 : index
+
+      %bcast = gpu.subgroup_broadcast %arg0, first_active_lane : memref<?xf32>
+      %bdim = memref.dim %bcast, %c0 : memref<?xf32>
+      %dim = memref.dim %arg0, %c0 : memref<?xf32>
+      // expected-remark @below{{true}}
+      "test.compare"(%bdim, %dim) {cmp = "EQ"} : (index, index) -> ()
+      gpu.return
+    }
+  }
+}