[mlir][gpu] Fix GPU YieldOP format and traits (#78006)

This patch adds assembly format to `gpu::YieldOp`. It also adds the
return like trait, to make it compatible with `RegionBranchOpInterface`.
This commit is contained in:
Fabian Mora 2024-01-14 21:19:20 -05:00 committed by GitHub
parent 2e0a105761
commit a1eaed7a21
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 17 additions and 1 deletions

View File

@ -23,6 +23,7 @@
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"

View File

@ -22,6 +22,7 @@ include "mlir/Dialect/GPU/TransformOps/GPUDeviceMappingAttr.td"
include "mlir/IR/CommonTypeConstraints.td"
include "mlir/IR/EnumAttr.td"
include "mlir/IR/SymbolInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/DataLayoutInterfaces.td"
include "mlir/Interfaces/FunctionInterfaces.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
@ -961,7 +962,7 @@ def GPU_TerminatorOp : GPU_Op<"terminator", [HasParent<"LaunchOp">,
let assemblyFormat = "attr-dict";
}
def GPU_YieldOp : GPU_Op<"yield", [Pure, Terminator]>,
def GPU_YieldOp : GPU_Op<"yield", [Pure, ReturnLike, Terminator]>,
Arguments<(ins Variadic<AnyType>:$values)> {
let summary = "GPU yield operation";
let description = [{
@ -974,6 +975,8 @@ def GPU_YieldOp : GPU_Op<"yield", [Pure, Terminator]>,
gpu.yield %f0, %f1 : f32, f32
```
}];
let assemblyFormat = "attr-dict ($values^ `:` type($values))?";
}
// These mirror the reduction combining kinds from the vector dialect.

View File

@ -37,6 +37,7 @@ add_mlir_dialect_library(MLIRGPUDialect
LINK_LIBS PUBLIC
MLIRArithDialect
MLIRDLTIDialect
MLIRControlFlowInterfaces
MLIRFunctionInterfaces
MLIRInferIntRangeInterface
MLIRIR

View File

@ -94,6 +94,17 @@ module attributes {gpu.container_module} {
// CHECK-NEXT: } : (f32) -> f32
%sum1 = gpu.all_reduce add %one uniform {} : (f32) -> f32
// CHECK: %{{.*}} = gpu.all_reduce %{{.*}} {
// CHECK-NEXT: ^{{.*}}(%{{.*}}: f32, %{{.*}}: f32):
// CHECK-NEXT: %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: gpu.yield %{{.*}} : f32
// CHECK-NEXT: } : (f32) -> f32
%sum2 = gpu.all_reduce %one {
^bb(%lhs : f32, %rhs : f32):
%tmp = arith.addf %lhs, %rhs : f32
gpu.yield %tmp : f32
} : (f32) -> (f32)
// CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} : (f32) -> f32
%sum_subgroup = gpu.subgroup_reduce add %one : (f32) -> f32