Follow up on custom instruction support.

This CL addresses a few post-submit comments:
1. better comments,
2. check number of results before dyn_cast (which is a less common case)
3. test usage for multi-result InstructionHandle

PiperOrigin-RevId: 237549333
This commit is contained in:
Nicolas Vasilache 2019-03-08 17:27:49 -08:00 committed by jpienaar
parent eb19b4eefc
commit 0d925c5510
3 changed files with 13 additions and 7 deletions

View File

@ -233,8 +233,8 @@ private:
BlockBuilder &operator=(const BlockBuilder &other) = delete;
};
/// Base class for Handles that cannot be constructed explicitly by a user of
/// the API.
/// Base class for ValueHandle, InstructionHandle and BlockHandle.
/// Not meant to be used outside of these classes.
struct CapturableHandle {
protected:
CapturableHandle() = default;
@ -347,6 +347,7 @@ struct InstructionHandle : public CapturableHandle {
ArrayRef<NamedAttribute> attributes = {});
operator Instruction *() { return inst; }
Instruction *getInstruction() { return inst; }
private:
Instruction *inst;

View File

@ -97,15 +97,15 @@ ValueHandle ValueHandle::create(StringRef name, ArrayRef<ValueHandle> operands,
ArrayRef<NamedAttribute> attributes) {
Instruction *inst =
InstructionHandle::create(name, operands, resultTypes, attributes);
if (inst->getNumResults() == 1) {
return ValueHandle(inst->getResult(0));
}
if (auto f = inst->dyn_cast<AffineForOp>()) {
// Immediately create the loop body so we can just insert instructions right
// away.
f->createBody();
return ValueHandle(f->getInductionVar());
}
if (inst->getNumResults() == 1) {
return ValueHandle(inst->getResult(0));
}
llvm_unreachable("unsupported instruction, use an InstructionHandle instead");
}

View File

@ -376,7 +376,7 @@ TEST_FUNC(custom_ops) {
CustomInstruction<InstructionHandle> MY_CUSTOM_INST_2("my_custom_inst_2");
// clang-format off
ValueHandle vh(indexType);
ValueHandle vh(indexType), vh20(indexType), vh21(indexType);
InstructionHandle ih0, ih2;
IndexHandle m, n, M(f->getArgument(0)), N(f->getArgument(1));
IndexHandle ten(index_t(10)), twenty(index_t(20));
@ -384,6 +384,10 @@ TEST_FUNC(custom_ops) {
vh = MY_CUSTOM_OP({m, m + n}, {indexType}, {}),
ih0 = MY_CUSTOM_INST_0({m, m + n}, {}),
ih2 = MY_CUSTOM_INST_2({m, m + n}, {indexType, indexType}),
// These captures are verbose for now, can improve when used in practice.
vh20 = ValueHandle(ih2.getInstruction()->getResult(0)),
vh21 = ValueHandle(ih2.getInstruction()->getResult(1)),
MY_CUSTOM_OP({vh20, vh21}, {indexType}, {}),
});
// CHECK-LABEL: @custom_ops
@ -391,7 +395,8 @@ TEST_FUNC(custom_ops) {
// CHECK: for %i1 {{.*}}
// CHECK: {{.*}} = "my_custom_op"{{.*}} : (index, index) -> index
// CHECK: "my_custom_inst_0"{{.*}} : (index, index) -> ()
// CHECK: {{.*}} = "my_custom_inst_2"{{.*}} : (index, index) -> (index, index)
// CHECK: [[TWO:%[a-z0-9]+]] = "my_custom_inst_2"{{.*}} : (index, index) -> (index, index)
// CHECK: {{.*}} = "my_custom_op"([[TWO]]#0, [[TWO]]#1) : (index, index) -> index
// clang-format on
f->print(llvm::outs());
}