[Mlir-commits] [mlir] 44027f3 - [mlir][GPU] Prevent adding duplicate async tokens

[Krzysztof Drewniak]

Author: Krzysztof Drewniak
Date: 2022-10-18T15:37:20Z
New Revision: 44027f390854f3088bf85c5c4d8cf1405dd4bb4c

URL: https://github.com/llvm/llvm-project/commit/44027f390854f3088bf85c5c4d8cf1405dd4bb4c
DIFF: https://github.com/llvm/llvm-project/commit/44027f390854f3088bf85c5c4d8cf1405dd4bb4c.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
    mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
    mlir/test/Dialect/GPU/async-region.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
index 31fc1683d683b..28bc3190c1450 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 a4feb9234e8ff..63eaa41fa0905 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 8bdafdafa9074..00832c45de114 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>