[mlir] Fix for gpu-async-region pass.

[csigg]

Summary

• the !gpu.async.token is the second result of 'gpu.alloc async', not the first.
• async.execute construction takes operand types not yet wrapped in !async.value.
• fix typo

[herhut]

There should be some test that captures this.

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

174

Should this be asserted in case the assumption changes?

[csigg]

In D93156#2451777, @herhut wrote:

There should be some test that captures this.

Done.

[csigg]

Thanks for the review!

Apply comments.

[herhut]

lgtm with extra test.

mlir/test/Dialect/GPU/async-region.mlir

30

Can you add an example where the async.execute returns a value beyond the token so that we can test that types get rewritten correctly?

[csigg]

Add test.

[csigg]

Rebase.

[csigg]

Rebase.

[csigg]

Trigger pre-merge checks.

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

LogicalResult rewriteAsyncOp(gpu::AsyncOpInterface asyncOp) {
 
     // If there is no current token, insert a `gpu.wait async` without
     // dependencies to create one.
     if (!currentToken)
       currentToken = createWaitOp(op->getLoc(), tokenType, {});
     asyncOp.addAsyncDependency(currentToken);
 
     // Clone the op to return a token in addition to the other results.
-    SmallVector<Type, 1> resultTypes = {tokenType};
+    SmallVector<Type, 1> resultTypes;
     resultTypes.reserve(1 + op->getNumResults());
     copy(op->getResultTypes(), std::back_inserter(resultTypes));
+    resultTypes.push_back(tokenType);
     auto *newOp = Operation::create(op->getLoc(), op->getName(), resultTypes,
                                     op->getOperands(), op->getMutableAttrDict(),
                                     op->getSuccessors());
 
     // Replace the op with the async clone.
     auto results = newOp->getResults();
-    currentToken = results.front();
+    currentToken = results.back();
     builder.insert(newOp);
-    op->replaceAllUsesWith(results.drop_front());
+    op->replaceAllUsesWith(results.drop_back());
     op->erase();
 
     return success();
   }
 
   Value createWaitOp(Location loc, Type resultType, ValueRange operands) {
     return builder.create<gpu::WaitOp>(loc, resultType, operands).asyncToken();
   }
private:
   // Append `count` `!async.value<!gpu.async.token>` results to `executeOp`.
   static ValueRange addAsyncTokenResults(async::ExecuteOp &executeOp,
                                          unsigned count) {
     auto numResults = executeOp.getNumResults() + count;
 
     // Construct new result type list with `count` additional types.
     SmallVector<Type, 2> resultTypes;
     resultTypes.reserve(numResults);
-    copy(executeOp.getResultTypes(), std::back_inserter(resultTypes));
+    transform(executeOp.getResultTypes(), std::back_inserter(resultTypes),
+              [](Type type) {
+                // Extract value type from !async.value.
+                if (auto valueType = type.dyn_cast<async::ValueType>())
+                  return valueType.getValueType();
+                assert(type.isa<async::TokenType>() && "expected token type");
+                return type;
+              });
     OpBuilder builder(executeOp);
     auto tokenType = builder.getType<gpu::AsyncTokenType>();
     resultTypes.resize(numResults, tokenType);
 
     // Clone executeOp with the extra `!gpu.async.token` results.
     auto newOp = builder.create<async::ExecuteOp>(
         executeOp.getLoc(), TypeRange{resultTypes}.drop_front() /*drop token*/,
         executeOp.dependencies(), executeOp.operands());

 // execution semantics and that no GPU ops are asynchronous yet.
 void GpuAsyncRegionPass::runOnFunction() {
   if (getFunction()
           .getRegion()
           .walk(ThreadTokenCallback(getContext()))
           .wasInterrupted())
     return signalPassFailure();
 
-  // Collect gpu.wait ops that we can move out of gpu.execute regions.
+  // Collect gpu.wait ops that we can move out of async.execute regions.
   getFunction().getRegion().walk(DeferWaitCallback());
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createGpuAsyncRegionPass() {
   return std::make_unique<GpuAsyncRegionPass>();
 }

mlir/test/Dialect/GPU/async-region.mlir

module attributes {gpu.container_module} {
   func @async(%sz : index) {
     // CHECK: %[[t0:.*]] = gpu.wait async
     // CHECK: %[[t1:.*]] = gpu.launch_func async [%[[t0]]]
     gpu.launch_func @kernels::@kernel
         blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
     // CHECK: %[[t2:.*]] = gpu.launch_func async [%[[t1]]]
     gpu.launch_func @kernels::@kernel
         blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
-    // CHECK: gpu.wait [%[[t2]]]
+    // CHECK: %[[m:.*]], %[[t3:.*]] = gpu.alloc async [%[[t2]]] ()
+    %0 = gpu.alloc() : memref<7xf32>
+    // CHECK: %[[t4:.*]] = gpu.dealloc async [%[[t3]]] %[[m]]
+    gpu.dealloc %0 : memref<7xf32>
+    // CHECK: gpu.wait [%[[t4]]]
     // CHECK: call @foo
     call @foo() : () -> ()
     return
   }
 
   // CHECK-LABEL:func @defer_wait(%{{.*}}: index)
   func @defer_wait(%sz : index) {
     // CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
     %a0 = async.execute {
       // CHECK: %[[t:.*]] = gpu.launch_func async
       gpu.launch_func @kernels::@kernel
           blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
       // CHECK-NOT: gpu.wait
       // CHECK: async.yield %[[t]]
       async.yield
     }
 
func @defer_wait_pass_through(%sz : index) {
     }
 
     // CHECK: async.await %[[a1]]
     // CHECK: %[[t:.*]] = async.await %[[f1]]
     // CHECK: gpu.wait [%[[t]]]
     async.await %a1 : !async.token
     return
   }
+
+ // CHECK-LABEL:func @async_execute_with_result(%{{.*}}: index)
+  func @async_execute_with_result(%sz : index) -> index {
+    // CHECK: %[[a0:.*]], %[[f0:.*]]:2 = async.execute
+    // CHECK-SAME: -> (!async.value<index>, !async.value<!gpu.async.token>)
+    %a0, %f0 = async.execute -> !async.value<index> {
+      // CHECK: %[[t:.*]] = gpu.launch_func async
+      gpu.launch_func @kernels::@kernel
+          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
+      // CHECK-NOT: gpu.wait
+      // CHECK: async.yield {{.*}}, %[[t]] : index, !gpu.async.token
+      async.yield %sz : index
+    }
+
+    // CHECK: async.await %[[a0]] : !async.token
+    // CHECK: %[[t:.*]] = async.await %[[f0]]#1 : !async.value<!gpu.async.token>
+    // CHECK: gpu.wait [%[[t]]]
+    async.await %a0 : !async.token
+    // CHECK: %[[x:.*]] = async.await %[[f0]]#0 : !async.value<index>
+    %x = async.await %f0 : !async.value<index>
+    // CHECK: return %[[x]] : index
+    return %x : index
+  }
 }