@@ -24,6 +24,7 @@ include "mlir/IR/EnumAttr.td"
24
24
include "mlir/IR/SymbolInterfaces.td"
25
25
include "mlir/Interfaces/ControlFlowInterfaces.td"
26
26
include "mlir/Interfaces/DataLayoutInterfaces.td"
27
+ include "mlir/IR/OpAsmInterface.td"
27
28
include "mlir/Interfaces/FunctionInterfaces.td"
28
29
include "mlir/Interfaces/InferIntRangeInterface.td"
29
30
include "mlir/Interfaces/InferTypeOpInterface.td"
@@ -50,9 +51,21 @@ def GPU_DimensionAttr : EnumAttr<GPU_Dialect, GPU_Dimension, "dim">;
50
51
51
52
class GPU_IndexOp<string mnemonic, list<Trait> traits = []> :
52
53
GPU_Op<mnemonic, !listconcat(traits, [
53
- Pure, DeclareOpInterfaceMethods<InferIntRangeInterface>])>,
54
+ Pure,
55
+ DeclareOpInterfaceMethods<InferIntRangeInterface>,
56
+ DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>])>,
54
57
Arguments<(ins GPU_DimensionAttr:$dimension)>, Results<(outs Index)> {
55
58
let assemblyFormat = "$dimension attr-dict";
59
+ let extraClassDefinition = [{
60
+ void $cppClass::getAsmResultNames(
61
+ llvm::function_ref<void(mlir::Value, mlir::StringRef)> setNameFn) {
62
+ auto dimStr = stringifyDimension(getDimensionAttr().getValue());
63
+ auto opName = getOperationName();
64
+ assert(opName.consume_front("gpu."));
65
+ SmallString<8> resultName({opName, "_", dimStr});
66
+ setNameFn(getResult(),resultName);
67
+ }
68
+ }];
56
69
}
57
70
58
71
def GPU_ClusterDimOp : GPU_IndexOp<"cluster_dim"> {
0 commit comments