[mlir] Prevent segfault in Tensor canonicalization

[tpopp]

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.

[mehdi_amini]

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]

Correct tests and fix comparison order.

[tpopp]

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]

Correct constant name

[tpopp]

Correct test case.

[rriddle]

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

253

Don't forget to remove this dump.

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

LogicalResult matchAndRewrite(tensor::ExtractOp extract,
 
     auto tensorFromElements = extract.tensor().getDefiningOp<FromElementsOp>();
     if (tensorFromElements == nullptr)
       return failure();
 
     APInt index;
     if (!matchPattern(*extract.indices().begin(), m_ConstantInt(&index)))
       return failure();
+    // Prevent out of bounds accesses. This can happen in invalid code that will
+    // never execute.
+    tensorFromElements.dump();
+    if (tensorFromElements->getNumOperands() <= index.getZExtValue() ||
+        index.getSExtValue() < 0)
+      return failure();
     rewriter.replaceOp(extract,
                        tensorFromElements.getOperand(index.getZExtValue()));
     return success();
   }
 };
 
 } // namespace
 

mlir/test/Dialect/Tensor/canonicalize.mlir

 // -----
 
 // CHECK-LABEL: func @extract_from_tensor.from_elements
 func @extract_from_tensor.from_elements(%element : index) -> index {
   // CHECK-SAME: ([[ARG:%.*]]: index)
   %c0 = constant 0 : index
   %tensor = tensor.from_elements %element : tensor<1xindex>
   %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex>
   // CHECK: [[ARG]] : 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 : tensor<0xindex>
+  %extracted_element = tensor.extract %tensor[%c1] : tensor<0xindex>
+  // 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-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32>
 func @extract_from_tensor.generate(%idx: index, %tensor: tensor<*xf32>) -> index {
   %size = rank %tensor : tensor<*xf32>
   // CHECK-NEXT: %[[RES:.*]] = dim %[[TENSOR]], %[[IDX]]
   %0 = tensor.generate %size {
     ^bb0(%arg0: index):
     %1 = dim %tensor, %arg0 : tensor<*xf32>