diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td @@ -724,6 +724,11 @@ in the shuffle. Exactly the first `width` invocations of a subgroup need to execute this op in convergence. }]; + let builders = [ + // Helper function that creates a shuffle with constant offset/width. + OpBuilder<(ins "Value":$value, "int32_t":$offset, "int32_t":$width, + "ShuffleMode":$mode)> + ]; let assemblyFormat = "$mode $value `,` $offset `,` $width attr-dict `:` type($value)"; } diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -691,6 +691,22 @@ printer << ")"; } +// + +//===----------------------------------------------------------------------===// +// ShuffleOp +//===----------------------------------------------------------------------===// + +void ShuffleOp::build(OpBuilder &builder, OperationState &result, Value value, + int32_t offset, int32_t width, ShuffleMode mode) { + build(builder, result, value, + builder.create(result.location, + builder.getI32IntegerAttr(offset)), + builder.create(result.location, + builder.getI32IntegerAttr(width)), + mode); +} + //===----------------------------------------------------------------------===// // GPUFuncOp //===----------------------------------------------------------------------===//