Skip to content

[mlir][IR] Support op interfaces in HasParent trait #91471

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][IR] Support op interfaces in HasParent trait #91471

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 1 addition & 2 deletions mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1463,8 +1463,7 @@ def Tensor_PadOp : Tensor_Op<"pad", [
def Tensor_ParallelInsertSliceOp : Tensor_Op<"parallel_insert_slice", [
AttrSizedOperandSegments,
OffsetSizeAndStrideOpInterface,
// TODO: Cannot use an interface here atm, verify this manually for now.
// HasParent<"ParallelCombiningOpInterface">
HasParent<"ParallelCombiningOpInterface">
]> {
let summary = [{
Specify the tensor slice update of a single thread of a parent
Expand Down
2 changes: 1 addition & 1 deletion mlir/include/mlir/IR/OpBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ class SingleBlockImplicitTerminator<string op>
// Op's regions don't have terminator.
def NoTerminator : NativeOpTrait<"NoTerminator">, StructuralOpTrait;

// Op's parent operation is the provided one.
// Op's parent operation or op interface is the provided one.
class HasParent<string op>
: ParamNativeOpTrait<"HasParent", op>, StructuralOpTrait;

Expand Down
23 changes: 22 additions & 1 deletion mlir/include/mlir/IR/OpDefinition.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1298,7 +1298,9 @@ struct HasParent {
return op->emitOpError()
<< "expects parent op "
<< (sizeof...(ParentOpTypes) != 1 ? "to be one of '" : "'")
<< llvm::ArrayRef({ParentOpTypes::getOperationName()...}) << "'";
<< llvm::ArrayRef(
{getOperationOrInterfaceName<ParentOpTypes>()...})
<< "'";
}

template <typename ParentOpType =
Expand All @@ -1309,6 +1311,25 @@ struct HasParent {
return llvm::cast<ParentOpType>(parent);
}
};

private:
/// A class is an op interface if it has a `getInterfaceName` function.
Copy link
Member Author

Choose a reason for hiding this comment

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

Ops and op interfaces inherit from OpState. If you know a better way of checking if a class is an interface, let me know. I also thought about calling the function getOperationName instead of getInterfaceName, but that could be confusing.

Copy link
Member

Choose a reason for hiding this comment

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

We actually have logic for this already at

llvm-project/mlir/include/mlir/Support/InterfaceSupport.h

Line 270 in 50b45b2

void isInterfaceImpl(
which is used to make DenseMap work with interfaces:

llvm-project/mlir/include/mlir/IR/OpDefinition.h

Line 2130 in 50b45b2

std::enable_if_t<std::is_base_of<mlir::OpState, T>::value &&

template <typename T, typename = int>
struct IsInterface : std::false_type {};
template <typename T>
struct IsInterface<T, decltype((void)T::getInterfaceName(), 0)>
: std::true_type {};

/// Helper function that returns the name of the given operation or interface
/// as a string literal.
template <typename T>
static constexpr StringLiteral getOperationOrInterfaceName() {
if constexpr (IsInterface<T>::value) {
return T::getInterfaceName();
} else {
return T::getOperationName();
}
}
};

/// A trait for operations that have an attribute specifying operand segments.
Expand Down
4 changes: 0 additions & 4 deletions mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3455,10 +3455,6 @@ void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result,
}

LogicalResult ParallelInsertSliceOp::verify() {
if (!isa<ParallelCombiningOpInterface>(getOperation()->getParentOp()))
return this->emitError("expected ParallelCombiningOpInterface parent, got:")
<< *(getOperation()->getParentOp());

RankedTensorType expectedType;
SliceVerificationResult result =
verifyInsertSliceOp(getSourceType(), getDestType(), getStaticOffsets(),
Expand Down
9 changes: 9 additions & 0 deletions mlir/test/Dialect/Tensor/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -698,3 +698,12 @@ func.func @unpack_mismatch_inner_tile_size_and_output_shape(
%0 = tensor.unpack %input inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %output : tensor<?x?x8x8xf32> -> tensor<?x?xf32>
return %0 : tensor<?x?xf32>
}

// -----

func.func @parallel_insert_slice_out_of_context(%a: tensor<5xf32>, %b: tensor<100xf32>) {
// expected-error@+1 {{expects parent op 'ParallelCombiningOpInterface'}}
tensor.parallel_insert_slice %a into %b[0][5][1]
: tensor<5xf32> into tensor<100xf32>
return
}
10 changes: 8 additions & 2 deletions mlir/tools/mlir-tblgen/OpInterfacesGen.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -537,7 +537,7 @@ void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {

// Emit the derived trait for the interface.
os << "template <typename " << valueTemplate << ">\n";
os << "struct " << interface.getName() << "Trait;\n";
os << "struct " << interfaceName << "Trait;\n";

os << "\n} // namespace detail\n";

Expand All @@ -548,6 +548,11 @@ void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {
interfaceName, interfaceName, interfaceTraitsName,
interfaceBaseType);

// Insert function that returns the name of the interface as a string.
os << " static constexpr ::llvm::StringLiteral getInterfaceName() {\n"
<< " return \"" << interfaceName << "\";\n"
<< " }\n\n";

// Emit a utility wrapper trait class.
os << llvm::formatv(" template <typename {1}>\n"
" struct Trait : public detail::{0}Trait<{1}> {{};\n",
Expand Down Expand Up @@ -588,7 +593,8 @@ void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {
<< " auto* interface = getInterfaceFor(base);\n"
<< " if (!interface)\n"
" return false;\n"
" " << interfaceName << " odsInterfaceInstance(base, interface);\n"
" "
<< interfaceName << " odsInterfaceInstance(base, interface);\n"
<< " " << tblgen::tgfmt(extraClassOf->trim(), &extraClassOfFmt)
<< "\n }\n";
}
Expand Down
Loading