Correctly model undefined behavior in {tensor|memref}.dim

[sanjoy]

Summary

These operations have undefined behavior if the index is not less than the rank of the source tensor / memref, so they cannot be freely speculated like they were before this patch. After this patch we speculate them only if we can prove that they don't have UB.

Depends on D135505.

[jpienaar]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

827

If you used ShapedType and hasRank then you'd be able to use same implementation for both ops.

[sanjoy]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

827

Specifically, are you suggesting I create a include/mlir/Dialect/Utils/DimOpUtils.h and put the shared implementation there?

[rriddle]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

827

I'd rather just duplicate the implementations than add more to Dialect/Utils.

833–834

Do you really need the (void) on these? The variables aren't unused AFAICT.

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

342

Same here.

mlir/test/Transforms/loop-invariant-code-motion.mlir

512

Please indent the checks with the IR.

[mravishankar]

Thanks Sanjoy! This makes sense to me.

[sanjoy]

Address review

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_DimOp : MemRef_Op<"dim", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     MemRefsNormalizable,
-    Pure,
+    ConditionallySpeculatable, NoMemoryEffect,
     ShapedDimOpInterface]> {
   let summary = "dimension index operation";
   let description = [{
     The `dim` operation takes a memref and a dimension operand of type `index`.
     It returns the size of the requested dimension of the given memref.
     If the dimension index is out of bounds the behavior is undefined.
 
     The specified memref type is that of the first operand.
let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
 
     /// Interface method of ShapedDimOpInterface: Return the source memref.
     Value getShapedValue() { return getSource(); }
 
     /// Interface method of ShapedDimOpInterface: Return the dimension.
     OpFoldResult getDimension() { return getIndex(); }
+
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

 }
 
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_DimOp : Tensor_Op<"dim", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
-    Pure,
+    ConditionallySpeculatable, NoMemoryEffect,
     ShapedDimOpInterface]> {
   let summary = "dimension index operation";
   let description = [{
     The `tensor.dim` operation takes a tensor and a dimension operand of type
     `index`. It returns the size of the requested dimension of the given
     tensor. If the dimension index is out of bounds, the behavior is undefined.
 
     The specified tensor type is that of the first operand.
let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
 
     /// Interface method of ShapedDimOpInterface: Return the source tensor.
     Value getShapedValue() { return getSource(); }
 
     /// Interface method of ShapedDimOpInterface: Return the dimension.
     OpFoldResult getDimension() { return getIndex(); }
+
+    /// Interface method for ConditionallySpeculatable.
+    Speculation::Speculatability getSpeculatability();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 }
 
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = getIndex().getDefiningOp<arith::ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 
+Speculation::Speculatability DimOp::getSpeculatability() {
+  auto constantIndex = getConstantIndex();
+  if (!constantIndex)
+    return Speculation::NotSpeculatable;
+
+  auto rankedSourceType = dyn_cast<MemRefType>(getSource().getType());
+  if (!rankedSourceType)
+    return Speculation::NotSpeculatable;
+
+  // The verifier rejects operations that violate this assertion.
+  assert(constantIndex < rankedSourceType.getRank());
+  (void) constantIndex;
+  (void) rankedSourceType;
+  return Speculation::Speculatable;
+}
+
 LogicalResult DimOp::verify() {
   // Assume unknown index to be in range.
   Optional<int64_t> index = getConstantIndex();
   if (!index)
     return success();
 
   // Check that constant index is not knowingly out of range.
   auto type = getSource().getType();

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

 }
 
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = getIndex().getDefiningOp<arith::ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 
+Speculation::Speculatability DimOp::getSpeculatability() {
+  auto constantIndex = getConstantIndex();
+  if (!constantIndex)
+    return Speculation::NotSpeculatable;
+
+  auto rankedSourceType = dyn_cast<RankedTensorType>(getSource().getType());
+  if (!rankedSourceType)
+    return Speculation::NotSpeculatable;
+
+  // The verifier rejects operations that violate this assertion.
+  assert(constantIndex < rankedSourceType.getRank());
+  (void) constantIndex;
+  (void) rankedSourceType;
+  return Speculation::Speculatable;
+}
+
 LogicalResult DimOp::verify() {
   // Assume unknown index to be in range.
   Optional<int64_t> index = getConstantIndex();
   if (!index)
     return success();
 
   // Check that constant index is not knowingly out of range.
   auto type = getSource().getType();

mlir/test/Transforms/loop-invariant-code-motion.mlir

scf.for %i = %lb to %ub step %step {
     %val = "test.recursively_speculatable_op"()({
       %result = "test.never_speculatable_op"() : () -> i32
       test.region_yield %result : i32
     }) : () -> i32
   }
 
   return
 }
+
+// -----
+
+func.func @speculate_tensor_dim_unknown_rank_unknown_dim(
+// CHECK-LABEL: @speculate_tensor_dim_unknown_rank_unknown_dim
+    %t: tensor<*xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+// CHECK: scf.for
+// CHECK-NEXT: tensor.dim
+  scf.for %i = %lb to %ub step %step {
+    %val = tensor.dim %t, %dim_idx : tensor<*xf32>
+  }
+
+  return
+}
+
+func.func @speculate_tensor_dim_known_rank_unknown_dim(
+// CHECK-LABEL: @speculate_tensor_dim_known_rank_unknown_dim
+    %t: tensor<?x?x?x?xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+// CHECK: scf.for
+// CHECK-NEXT: tensor.dim
+  scf.for %i = %lb to %ub step %step {
+    %val = tensor.dim %t, %dim_idx : tensor<?x?x?x?xf32>
+  }
+
+  return
+}
+
+func.func @speculate_tensor_dim_unknown_rank_known_dim(
+// CHECK-LABEL: @speculate_tensor_dim_unknown_rank_known_dim
+    %t: tensor<*xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+  %c0 = arith.constant 0 : index
+// CHECK: scf.for
+// CHECK-NEXT: tensor.dim
+  scf.for %i = %lb to %ub step %step {
+    %val = tensor.dim %t, %c0 : tensor<*xf32>
+  }
+
+  return
+}
+
+func.func @speculate_tensor_dim_known_rank_known_dim_inbounds(
+// CHECK-LABEL: @speculate_tensor_dim_known_rank_known_dim_inbounds
+    %t: tensor<?x?x?x?xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+  %c1 = arith.constant 1 : index
+// CHECK: tensor.dim
+// CHECK-NEXT: scf.for
+  scf.for %i = %lb to %ub step %step {
+    %val = tensor.dim %t, %c1 : tensor<?x?x?x?xf32>
+  }
+
+  return
+}
+
+// -----
+
+func.func @speculate_memref_dim_unknown_rank_unknown_dim(
+// CHECK-LABEL: @speculate_memref_dim_unknown_rank_unknown_dim
+    %t: memref<*xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+// CHECK: scf.for
+// CHECK-NEXT: memref.dim
+  scf.for %i = %lb to %ub step %step {
+    %val = memref.dim %t, %dim_idx : memref<*xf32>
+  }
+
+  return
+}
+
+func.func @speculate_memref_dim_known_rank_unknown_dim(
+// CHECK-LABEL: @speculate_memref_dim_known_rank_unknown_dim
+    %t: memref<?x?x?x?xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+// CHECK: scf.for
+// CHECK-NEXT: memref.dim
+  scf.for %i = %lb to %ub step %step {
+    %val = memref.dim %t, %dim_idx : memref<?x?x?x?xf32>
+  }
+
+  return
+}
+
+func.func @speculate_memref_dim_unknown_rank_known_dim(
+// CHECK-LABEL: @speculate_memref_dim_unknown_rank_known_dim
+    %t: memref<*xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+  %c0 = arith.constant 0 : index
+// CHECK: scf.for
+// CHECK-NEXT: memref.dim
+  scf.for %i = %lb to %ub step %step {
+    %val = memref.dim %t, %c0 : memref<*xf32>
+  }
+
+  return
+}
+
+func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
+// CHECK-LABEL: @speculate_memref_dim_known_rank_known_dim_inbounds
+    %t: memref<?x?x?x?xf32>, %dim_idx: index, %lb: index, %ub: index, %step: index) {
+  %c1 = arith.constant 1 : index
+// CHECK: memref.dim
+// CHECK-NEXT: scf.for
+  scf.for %i = %lb to %ub step %step {
+    %val = memref.dim %t, %c1 : memref<?x?x?x?xf32>
+  }
+
+  return
+}