This commit is contained in:
@@ -6,6 +6,7 @@ include "mlir/IR/OpAsmInterface.td"
|
||||
include "mlir/IR/BuiltinTypes.td"
|
||||
include "mlir/IR/AttrTypeBase.td"
|
||||
include "mlir/IR/RegionKindInterface.td"
|
||||
include "mlir/Interfaces/ControlFlowInterfaces.td"
|
||||
include "mlir/Interfaces/ParallelCombiningOpInterface.td"
|
||||
include "mlir/Interfaces/SideEffectInterfaces.td"
|
||||
|
||||
@@ -27,7 +28,7 @@ def SpatTensor :
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
|
||||
[SingleBlock, AttrSizedOperandSegments,
|
||||
[AttrSizedOperandSegments,
|
||||
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
|
||||
let summary = "Compute region with attached constant weights";
|
||||
|
||||
@@ -40,7 +41,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
|
||||
Variadic<SpatTensor>:$outputs
|
||||
);
|
||||
|
||||
let regions = (region SizedRegion<1>:$body);
|
||||
let regions = (region MinSizedRegion<1>:$body);
|
||||
|
||||
let hasVerifier = 1;
|
||||
let hasFolder = 1;
|
||||
@@ -76,7 +77,7 @@ def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> {
|
||||
}
|
||||
|
||||
class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
|
||||
[SingleBlock, AttrSizedOperandSegments,
|
||||
[AttrSizedOperandSegments,
|
||||
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
|
||||
let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";
|
||||
|
||||
@@ -90,7 +91,7 @@ class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
|
||||
Variadic<SpatTensor>:$outputs
|
||||
);
|
||||
|
||||
let regions = (region SizedRegion<1>:$body);
|
||||
let regions = (region MinSizedRegion<1>:$body);
|
||||
|
||||
let hasVerifier = 1;
|
||||
let hasCustomAssemblyFormat = 1;
|
||||
@@ -161,6 +162,54 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> {
|
||||
let hasCustomAssemblyFormat = 1;
|
||||
}
|
||||
|
||||
def SpatBlockYieldOp : SpatOp<"block_yield", [
|
||||
Terminator,
|
||||
DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>
|
||||
]> {
|
||||
let summary = "Terminate a scheduled structural compute block";
|
||||
|
||||
let arguments = (ins
|
||||
Variadic<AnyType>:$outputs
|
||||
);
|
||||
|
||||
let successors = (successor
|
||||
VariadicSuccessor<AnySuccessor>:$next
|
||||
);
|
||||
|
||||
let hasVerifier = 1;
|
||||
let hasCustomAssemblyFormat = 1;
|
||||
}
|
||||
|
||||
def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> {
|
||||
let summary = "Temporary scheduled communication placeholder";
|
||||
|
||||
let arguments = (ins
|
||||
Variadic<SpatTensor>:$sources,
|
||||
StrAttr:$exchangeId,
|
||||
StrAttr:$logicalProducer,
|
||||
StrAttr:$logicalConsumer,
|
||||
I64Attr:$sourceClass,
|
||||
I64Attr:$targetClass,
|
||||
I64Attr:$sourceCore,
|
||||
I64Attr:$targetCore,
|
||||
I64Attr:$sourceLane,
|
||||
I64Attr:$targetLane,
|
||||
StrAttr:$transferKind,
|
||||
I64Attr:$resultIndex,
|
||||
DenseI64ArrayAttr:$projectedTransfer,
|
||||
I64Attr:$hostOutputOwner
|
||||
);
|
||||
|
||||
let results = (outs
|
||||
SpatTensor:$output
|
||||
);
|
||||
|
||||
let regions = (region SizedRegion<1>:$body);
|
||||
|
||||
let hasVerifier = 1;
|
||||
let hasCustomAssemblyFormat = 1;
|
||||
}
|
||||
|
||||
def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
|
||||
let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
|
||||
|
||||
|
||||
Reference in New Issue
Block a user