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

[Jakub Kuderski]

================
@@ -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)";
+}
----------------
kuhar wrote:

Another similar op is `util.assume` in IREE which also takes a value and returns a value, unlike [llvm.assume](https://llvm.org/docs/LangRef.html#llvm-assume-intrinsic). We found that it's much easier to follow SSA use-def chains to discover assumptions than to have some separate assumption cache on the side.

Semantically, the only thing these ops do is to add more constraints to the path condition, but do not change the execution semantics. I think that's perfectly reasonable.

RE 3:
> In %v2, is %v0 uniform?

an analysis could conclude both are uniform

> Are %v2 and %v3 the same value?

an analysis could conclude they have the same value

> Can one be marked as uniform and the other not?

yep

> These can be defined, but I'd rather avoid defining what happens in those cases.

I think the current op semantics are pretty clear when it comes to this. Is there something that would not allow you to answer these questions?

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