This is an archive of the discontinued LLVM Phabricator instance.

Differential D95535

[mlir] Prevent segfault in Tensor canonicalization
ClosedPublic

Authored by tpopp on Jan 27 2021, 8:30 AM.

Details

Reviewers
silvas
mehdi_amini
Commits
rG0c5e4a25ee23: [mlir] Prevent segfault in Tensor canonicalization
Summary

This segfault could occur from out of bounds accesses when simplifying
tensor.extract with a constant index and a tensor created by
tensor.from_elements.

This IR is not necesarilly invalid as it might conditionally be
never executed.

Diff Detail

Event Timeline

tpopp created this revision.Jan 27 2021, 8:30 AM
Herald added subscribers: teijeong, rdzhabarov, tatianashp and 14 others. · View Herald TranscriptJan 27 2021, 8:30 AM
tpopp requested review of this revision.Jan 27 2021, 8:30 AM
Herald added a project: Restricted Project. · View Herald TranscriptJan 27 2021, 8:30 AM
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
Harbormaster completed remote builds in B86856: Diff 319589.Jan 27 2021, 8:59 AM
mehdi_amini added inline comments.Jan 27 2021, 5:54 PM
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
253

I don't understand this check, it reads backward to me right now?

mlir/test/Dialect/Tensor/canonicalize.mlir
119

Why is this check removed?

tpopp updated this revision to Diff 319798.Jan 28 2021, 2:16 AM

Correct tests and fix comparison order.

tpopp marked 2 inline comments as done.Jan 28 2021, 2:19 AM
tpopp added inline comments.
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
253

It was backwards, and my test only involved the = case. I've now added another for the < case.

mlir/test/Dialect/Tensor/canonicalize.mlir
119

Thanks for catching that.. Added back.

tpopp updated this revision to Diff 319800.Jan 28 2021, 2:20 AM
tpopp marked 2 inline comments as done.

Correct constant name

Harbormaster completed remote builds in B86987: Diff 319798.Jan 28 2021, 2:53 AM
tpopp updated this revision to Diff 319817.Jan 28 2021, 3:14 AM

Correct test case.

Harbormaster completed remote builds in B86989: Diff 319800.Jan 28 2021, 3:16 AM
Harbormaster completed remote builds in B86997: Diff 319817.Jan 28 2021, 3:43 AM
mehdi_amini accepted this revision.Jan 28 2021, 4:10 PM
This revision is now accepted and ready to land.Jan 28 2021, 4:10 PM
rriddle added inline comments.Jan 28 2021, 4:18 PM
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
253

Don't forget to remove this dump.

This revision was landed with ongoing or failed builds.Jan 29 2021, 1:58 AM
Closed by commit rG0c5e4a25ee23: [mlir] Prevent segfault in Tensor canonicalization (authored by Tres Popp <tpopp@google.com>). · Explain Why
This revision was automatically updated to reflect the committed changes.
Tres Popp <tpopp@google.com> added a commit: rG0c5e4a25ee23: [mlir] Prevent segfault in Tensor canonicalization.

Revision Contents

PathSize
mlir/
lib/
Dialect/
Tensor/
IR/
5 lines
test/
Dialect/
Tensor/
45 lines

Diff 320076

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Show First 20 Lines • Show All 242 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tensor::ExtractOp extract,
auto tensorFromElements = extract.tensor().getDefiningOp<FromElementsOp>(); auto tensorFromElements = extract.tensor().getDefiningOp<FromElementsOp>();
if (tensorFromElements == nullptr) if (tensorFromElements == nullptr)
return failure(); return failure();
APInt index; APInt index;
if (!matchPattern(*extract.indices().begin(), m_ConstantInt(&index))) if (!matchPattern(*extract.indices().begin(), m_ConstantInt(&index)))
return failure(); return failure();
// Prevent out of bounds accesses. This can happen in invalid code that will
// never execute.
if (tensorFromElements->getNumOperands() <= index.getZExtValue() ||
mehdi_aminiUnsubmitted
Done
ReplyInline Actions

I don't understand this check, it reads backward to me right now?

mehdi_amini: I don't understand this check, it reads backward to me right now?
tpoppAuthorUnsubmitted
Done
ReplyInline Actions

It was backwards, and my test only involved the = case. I've now added another for the < case.

tpopp: It was backwards, and my test only involved the `=` case. I've now added another for the `<`…
rriddleUnsubmitted
Not Done
ReplyInline Actions

Don't forget to remove this dump.

rriddle: Don't forget to remove this dump.
index.getSExtValue() < 0)
return failure();
rewriter.replaceOp(extract, rewriter.replaceOp(extract,
tensorFromElements.getOperand(index.getZExtValue())); tensorFromElements.getOperand(index.getZExtValue()));
return success(); return success();
} }
}; };
} // namespace } // namespace
▲ Show 20 Lines • Show All 173 Lines • Show Last 20 Lines

mlir/test/Dialect/Tensor/canonicalize.mlir

Show First 20 Lines • Show All 110 Lines • ▼ Show 20 Lines
// ----- // -----
// CHECK-LABEL: func @extract_from_tensor.from_elements // CHECK-LABEL: func @extract_from_tensor.from_elements
func @extract_from_tensor.from_elements(%element : index) -> index { func @extract_from_tensor.from_elements(%element : index) -> index {
// CHECK-SAME: ([[ARG:%.*]]: index) // CHECK-SAME: ([[ARG:%.*]]: index)
%c0 = constant 0 : index %c0 = constant 0 : index
%tensor = tensor.from_elements %element : tensor<1xindex> %tensor = tensor.from_elements %element : tensor<1xindex>
%extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex> %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex>
// CHECK: [[ARG]] : index // CHECK: [[ARG]] : index
mehdi_aminiUnsubmitted
Done
ReplyInline Actions

Why is this check removed?

mehdi_amini: Why is this check removed?
tpoppAuthorUnsubmitted
Done
ReplyInline Actions

Thanks for catching that.. Added back.

tpopp: Thanks for catching that.. Added back.
return %extracted_element : index return %extracted_element : index
} }
// ----- // -----
// Ensure the optimization doesn't segfault from bad constants
// CHECK-LABEL: func @extract_negative_from_tensor.from_elements
func @extract_negative_from_tensor.from_elements(%element : index) -> index {
// CHECK-SAME: ([[ARG:%.*]]: index)
%c-1 = constant -1 : index
%tensor = tensor.from_elements %element : tensor<1xindex>
%extracted_element = tensor.extract %tensor[%c-1] : tensor<1xindex>
// CHECK: tensor.from_elements
// CHECK: %[[RESULT:.*]] = tensor.extract
// CHECK: return %[[RESULT]]
return %extracted_element : index
}
// -----
// Ensure the optimization doesn't segfault from bad constants
// CHECK-LABEL: func @extract_oob_from_tensor.from_elements
func @extract_oob_from_tensor.from_elements(%element : index) -> index {
// CHECK-SAME: ([[ARG:%.*]]: index)
%c1 = constant 1 : index
%tensor = tensor.from_elements %element : tensor<1xindex>
%extracted_element = tensor.extract %tensor[%c1] : tensor<1xindex>
// CHECK: tensor.from_elements
// CHECK: %[[RESULT:.*]] = tensor.extract
// CHECK: return %[[RESULT]]
return %extracted_element : index
}
// -----
// Ensure the optimization doesn't segfault from bad constants
// CHECK-LABEL: func @extract_oob_from_tensor.from_elements
func @extract_oob_from_tensor.from_elements(%element : index) -> index {
// CHECK-SAME: ([[ARG:%.*]]: index)
%c2 = constant 2 : index
%tensor = tensor.from_elements %element : tensor<1xindex>
%extracted_element = tensor.extract %tensor[%c2] : tensor<1xindex>
// CHECK: tensor.from_elements
// CHECK: %[[RESULT:.*]] = tensor.extract
// CHECK: return %[[RESULT]]
return %extracted_element : index
}
// -----
// CHECK-LABEL: func @extract_from_tensor.generate // CHECK-LABEL: func @extract_from_tensor.generate
// CHECK-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32> // CHECK-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32>
func @extract_from_tensor.generate(%idx: index, %tensor: tensor<*xf32>) -> index { func @extract_from_tensor.generate(%idx: index, %tensor: tensor<*xf32>) -> index {
%size = rank %tensor : tensor<*xf32> %size = rank %tensor : tensor<*xf32>
// CHECK-NEXT: %[[RES:.*]] = dim %[[TENSOR]], %[[IDX]] // CHECK-NEXT: %[[RES:.*]] = dim %[[TENSOR]], %[[IDX]]
%0 = tensor.generate %size { %0 = tensor.generate %size {
^bb0(%arg0: index): ^bb0(%arg0: index):
%1 = dim %tensor, %arg0 : tensor<*xf32> %1 = dim %tensor, %arg0 : tensor<*xf32>
▲ Show 20 Lines • Show All 64 Lines • Show Last 20 Lines