Skip to content

[Dialect] [Linalgx] Add linalgx ops: 3 vnni matmuls and multi_batch_matmul #89

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 19 commits into from
Jun 3, 2024
Merged

[Dialect] [Linalgx] Add linalgx ops: 3 vnni matmuls and multi_batch_matmul #89

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
959eb93
add basic linalgx
May 20, 2024
239d1ac
Merge branch 'main' into longsheng/add_linalgx
May 20, 2024
5e88a8c
merge
May 20, 2024
9bae572
fix
May 20, 2024
3ac1a40
add
May 20, 2024
2cd9906
action: generate multiple inc target in one makefile
WangJialei-A May 20, 2024
6022f40
fix
May 20, 2024
284f83d
Merge branch 'longsheng/add_linalgx' into longsheng/add_linalgx_mm
May 20, 2024
c397b88
update
May 20, 2024
4ae4f17
update
May 20, 2024
e192a7a
update
May 21, 2024
a032b70
fix mmt2d
May 21, 2024
c8e31ba
Merge branch 'main' into longsheng/add_linalgx
May 21, 2024
546fccb
Merge branch 'longsheng/add_linalgx' into longsheng/add_linalgx_mm
May 21, 2024
5529faf
add brgemm vnni
May 22, 2024
5bf2b66
Merge branch 'main' into longsheng/add_linalgx_mm
Jun 3, 2024
d97b267
update
Jun 3, 2024
fcbb679
fix test
Jun 3, 2024
76e6291
fix const affine_map
Jun 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 208 additions & 0 deletions include/gc/Dialect/Linalgx/LinalgxStructuredOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,4 +104,212 @@ def Linalgx_SigmoidOp : LinalgxStructuredBase_Op<"sigmoid",
}];
}

def Linalgx_Mm2DVnniOp
: LinalgxStructuredBase_Op<"mm2d_vnni", [AttrSizedOperandSegments]> {
let summary = "Transposed matmul with 2d input and vnni packed weights";
let description = [{
Supported format: A[M, K] * B[N0, K0, k, n, v] -> C[M, N], with:
N = N0 * n
K = K0 * k * v; v = (2, 4)
}];
let arguments = (ins
Variadic<TensorOrMemref>:$inputs,
Variadic<TensorOrMemref>:$outputs);
let results = (outs Variadic<TensorOrMemref>:$results);
let regions = (region AnyRegion:$region);

let skipDefaultBuilders = 1;
let builders = [
OpBuilder<
(ins
"TypeRange":$resultTensorTypes,
"ValueRange":$inputs,
"ValueRange":$outputs,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
[{
buildStructuredOp($_builder, $_state, resultTensorTypes,
inputs, outputs, attributes, Mm2DVnniOp::getRegionBuilder());
}]>
];

let hasCustomAssemblyFormat = 1;
let hasFolder = 1;
let hasVerifier = 1;

let extraClassDeclaration = structuredOpsBaseDecls # [{
// Declare functions necessary for LinalgStructuredInterface.
SmallVector<utils::IteratorType> getIteratorTypesArray();
ArrayAttr getIndexingMaps();
static unsigned getNumRegionArgs() { return 3; }
std::string getLibraryCallName() {
return "op_has_no_registered_library_name";
}

// Implement functions necessary for DestinationStyleOpInterface.
MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }

static void regionBuilder(ImplicitLocOpBuilder &b,
Block &block, ArrayRef<NamedAttribute> attrs);
static std::function<void(ImplicitLocOpBuilder &,
Block &, ArrayRef<NamedAttribute>)>
getRegionBuilder() {
return regionBuilder;
}
}];
}

def Linalgx_Mm4DVnniOp
: LinalgxStructuredBase_Op<"mm4d_vnni", [AttrSizedOperandSegments]> {
let summary = "Transposed matmul with 4d blocking input and vnni packed weights";
let description = [{
Supported format: A[M, K, m, k] * B[N, K, k0, n, v] -> C[M, N, m, n], with:
k = k0 * v; v = (2, 4)
}];
let arguments = (ins
Variadic<TensorOrMemref>:$inputs,
Variadic<TensorOrMemref>:$outputs);
let results = (outs Variadic<TensorOrMemref>:$results);
let regions = (region AnyRegion:$region);

let skipDefaultBuilders = 1;
let builders = [
OpBuilder<
(ins
"TypeRange":$resultTensorTypes,
"ValueRange":$inputs,
"ValueRange":$outputs,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
[{
buildStructuredOp($_builder, $_state, resultTensorTypes,
inputs, outputs, attributes, Mm4DVnniOp::getRegionBuilder());
}]>
];

let hasCustomAssemblyFormat = 1;
let hasFolder = 1;
let hasVerifier = 1;

let extraClassDeclaration = structuredOpsBaseDecls # [{
// Declare functions necessary for LinalgStructuredInterface.
SmallVector<utils::IteratorType> getIteratorTypesArray();
ArrayAttr getIndexingMaps();
static unsigned getNumRegionArgs() { return 3; }
std::string getLibraryCallName() {
return "op_has_no_registered_library_name";
}

// Implement functions necessary for DestinationStyleOpInterface.
MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }

static void regionBuilder(ImplicitLocOpBuilder &b,
Block &block, ArrayRef<NamedAttribute> attrs);
static std::function<void(ImplicitLocOpBuilder &,
Block &, ArrayRef<NamedAttribute>)>
getRegionBuilder() {
return regionBuilder;
}
}];
}

def Linalgx_BatchReduceMatmulVnniOp
: LinalgxStructuredBase_Op<"batch_reduce_matmul_vnni", [AttrSizedOperandSegments]> {
let summary = "Batch reduced matmul with 3d batch input and vnni packed weights";
let description = [{
Supported format: A[B, M, K] * B[B, k, N, v] -> C[M, N], with:
K = k * v; v = (2, 4)
}];
let arguments = (ins
Variadic<TensorOrMemref>:$inputs,
Variadic<TensorOrMemref>:$outputs);
let results = (outs Variadic<TensorOrMemref>:$results);
let regions = (region AnyRegion:$region);

let skipDefaultBuilders = 1;
let builders = [
OpBuilder<
(ins
"TypeRange":$resultTensorTypes,
"ValueRange":$inputs,
"ValueRange":$outputs,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
[{
buildStructuredOp($_builder, $_state, resultTensorTypes,
inputs, outputs, attributes, BatchReduceMatmulVnniOp::getRegionBuilder());
}]>
];

let hasCustomAssemblyFormat = 1;
let hasFolder = 1;
let hasVerifier = 1;

let extraClassDeclaration = structuredOpsBaseDecls # [{
// Declare functions necessary for LinalgStructuredInterface.
SmallVector<utils::IteratorType> getIteratorTypesArray();
ArrayAttr getIndexingMaps();
static unsigned getNumRegionArgs() { return 3; }
std::string getLibraryCallName() {
return "op_has_no_registered_library_name";
}

// Implement functions necessary for DestinationStyleOpInterface.
MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }

static void regionBuilder(ImplicitLocOpBuilder &b,
Block &block, ArrayRef<NamedAttribute> attrs);
static std::function<void(ImplicitLocOpBuilder &,
Block &, ArrayRef<NamedAttribute>)>
getRegionBuilder() {
return regionBuilder;
}
}];
}

def Linalgx_MultiBatchMatmulOp : LinalgxStructuredBase_Op<"multi_batch_matmul",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also define multi_batch_matmul_4d and multi_batch_matmul_4d_vnni?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this define correct for multi_batch_matmul_4d and multi_batch_matmul_4d_vnni? Do we also want transposed weight?

input: BMKmk
weight: BNKkn4k/BNKkn2k/BNKkn1k
output: BMNmn

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct. We don't use transposed weight.

[AttrSizedOperandSegments, LinalgContractionOpInterface]> {
let summary = "Batch matmul with variable batch dims";
let arguments = (ins
Variadic<TensorOrMemref>:$inputs,
Variadic<TensorOrMemref>:$outputs);
let results = (outs Variadic<TensorOrMemref>:$results);
let regions = (region AnyRegion:$region);

let skipDefaultBuilders = 1;
let builders = [
OpBuilder<
(ins
"TypeRange":$resultTensorTypes,
"ValueRange":$inputs,
"ValueRange":$outputs,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
[{
buildStructuredOp($_builder, $_state, resultTensorTypes,
inputs, outputs, attributes, MultiBatchMatmulOp::getRegionBuilder());
}]>
];

let hasCustomAssemblyFormat = 1;
let hasFolder = 1;

let extraClassDeclaration = structuredOpsBaseDecls # [{
// Declare functions necessary for LinalgStructuredInterface.
SmallVector<utils::IteratorType> getIteratorTypesArray();
ArrayAttr getIndexingMaps();
static unsigned getNumRegionArgs() { return 3; }
std::string getLibraryCallName() {
return "op_has_no_registered_library_name";
}

// Implement functions necessary for DestinationStyleOpInterface.
MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }

static void regionBuilder(ImplicitLocOpBuilder &b,
Block &block, ArrayRef<NamedAttribute> attrs);
static std::function<void(ImplicitLocOpBuilder &,
Block &, ArrayRef<NamedAttribute>)>
getRegionBuilder() {
return regionBuilder;
}
}];
}

#endif // LINALGX_STRUCTURED_OPS
Loading