[Mlir-commits] [mlir] [mlir][gpu] Add `gpu.subgroup_uniform` op (PR #157743)

[Krzysztof Drewniak]

================
@@ -3255,4 +3255,37 @@ def GPU_SubgroupBroadcastOp : GPU_Op<"subgroup_broadcast",
   let hasVerifier = 1;
 }
 
+def GPU_SubgroupUniformOp : GPU_Op<"subgroup_uniform",
+    [Pure, AllTypesMatch<["result", "src"]>,
+    DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>] #
+    ElementwiseMappable.traits>,
+  Arguments<(ins AnyType:$src)> {
+  let summary = "Assumes value is unform across the lanes in subgroup";
+  let description = [{
+    The "subgroup_uniform" op assumes that the value is uniform across all lanes
+    in a subgroup. This means that all active lanes in the subgroup are expected
+    to have the same value.
+
+    This op can be used to inform the compiler that a value is uniform across
+    the subgroup, enabling optimizations. The result is poison if the value
+    is not actually uniform.
+
+    This op is functionally no-op as no valid program should change its
+    semantics if this op is removed. Backends can choose to ignore it or do
+    some optimizations (e.g. put value into scalar registers).
+
+    This op can be freely speculated across structured control flow as parent
+    active mask is always superset of current mask and if can hoist input
+    calculation you can hoist the operation itself as well.
+
+    Example:
+
+    ```mlir
+    %1 = gpu.subgroup_uniform %0 : f32
+    ```
+  }];
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "$src attr-dict `:` type($result)";
+}
----------------
krzysz00 wrote:

@fabianmcg  And re the `my_special_op` example - can you dream up a `my_special_op` that would cause `%v0` to not be subgroup-uniform outside of `my_special_op` but uniform inside of it?

Otherwise, I claim the following rewrite is correct

```
%v0 = // ....
%u0 = gpu.subgroup_uniform %v0 : ...
my.special_op {
  // ...
}
// [tag] %v0 use
```

But, outside of hoisting our way out of this situation ... a thing that analyzes uniformity should not use the fact that `%u0` exists to draw conclusions about `%v0`.

That is, the assumption (and any actual concrete implemetation effects) flows from `%u0` to its users, and `%v0` isn't affected by them.

(See also the `assume_alignment` example. You only get the assumption on the users of the result of `assume_alignment`.)

Now, is it legal to replaceAllUsesWith(%v0, %u0) ... well, probably not for SSA reasons, but if it weren't for the control flow - if this were straight-line code - it would've been.

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