[mlir][GPU] Prevent adding duplicate async tokens

If, in the GPU async transformation, the operation being given an
async dependency already depended on the token in question, we
would add duplicate tokens, creating issues in GPU to LLVM lowering.

To resolve this issue, add a check to addAsyncDependency() to ensure
that duplicate tokens are not present in the token list.

(I'm open to a different approach here, this is just what I went with
initially)

Reviewed By: Mogball

Differential Revision: https://reviews.llvm.org/D136105

diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
index 31fc168..28bc319 100644 (file)
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
@@ -96,15 +96,16 @@ def GPU_AsyncOpInterface : OpInterface<"AsyncOpInterface"> {
       }],
       "OperandRange", "getAsyncDependencies", (ins), [{}], [{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
-        return op.asyncDependencies();
+        return op.getAsyncDependencies();
       }]
     >,
     InterfaceMethod<[{
-        Adds a new token to the list of async dependencies.
+        Adds a new token to the list of async dependencies if it is not already there.
       }],
       "void", "addAsyncDependency", (ins "Value":$token),
       [{}], [{
-        ::mlir::gpu::addAsyncDependency(this->getOperation(), token);
+        if (!::llvm::is_contained(this->getAsyncDependencies(), token))
+          ::mlir::gpu::addAsyncDependency(this->getOperation(), token);
       }]
     >,
     InterfaceMethod<[{

diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
index a4feb92..63eaa41 100644 (file)
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
@@ -25,6 +25,7 @@
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace mlir {
 namespace gpu {

diff --git a/mlir/test/Dialect/GPU/async-region.mlir b/mlir/test/Dialect/GPU/async-region.mlir
index 8bdafda..00832c4 100644 (file)
--- a/mlir/test/Dialect/GPU/async-region.mlir
+++ b/mlir/test/Dialect/GPU/async-region.mlir
@@ -175,7 +175,7 @@ module attributes {gpu.container_module} {
     // CHECK: %[[t0:.*]] = gpu.wait async
     // CHECK-NOT: [{{.*}}]
     %t0 = gpu.wait async
-    // CHECK: %[[t1:.*]] = gpu.wait async [%[[t0]], %[[t0]]]
+    // CHECK: %[[t1:.*]] = gpu.wait async [%[[t0]]]
     %t1 = gpu.wait async [%t0]
     // CHECK: %[[m:.*]], %[[t2:.*]] = gpu.alloc async [%[[t1]], %[[t0]]] ()
     %0 = gpu.alloc [%t0] () : memref<7xf32>