From 44027f390854f3088bf85c5c4d8cf1405dd4bb4c Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak@amd.com>
Date: Mon, 17 Oct 2022 17:47:56 +0000
Subject: [PATCH] [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
---
 mlir/include/mlir/Dialect/GPU/IR/GPUBase.td   | 7 ++++---
 mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h | 1 +
 mlir/test/Dialect/GPU/async-region.mlir       | 2 +-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
index 31fc168..28bc319 100644
--- 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
--- 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
--- 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>
-- 
2.7.4