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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Krzysztof Drewniak (krzysz00)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/183848.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp (+23) 
- (modified) mlir/test/Dialect/GPU/value-bounds-op-interface-impl.mlir (+38) 


``````````diff
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
+    }
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/183848