[mlir][MemRef] Add ValueBoundsOpInterface impls

[springerm]

Summary

Depends On: D145694

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

update API

[dcaballe]

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

108

-> i =0, end = getType().getRank()

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

29 ↗

(On Diff #509964)

ok, this seems to be common so I'm guessing the operator == is overloaded?

[springerm]

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

29 ↗

(On Diff #509964)

Yes, this is the new API that jpienaar suggested.

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

OpBuilder<(ins "MemRefType":$memrefType, "ValueRange":$dynamicSizes,
       if (alignment)
         $_state.addAttribute(getAlignmentAttrStrName(), alignment);
     }]>];
 
   let extraClassDeclaration = [{
     static StringRef getAlignmentAttrStrName() { return "alignment"; }
 
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
+
+    SmallVector<OpFoldResult> getMixedSizes() {
+      SmallVector<OpFoldResult> result;
+      unsigned ctr = 0;
+      OpBuilder b(getContext());
+      for (int64_t i = 0; i < getType().getRank(); ++i) {
+        if (getType().isDynamicDim(i)) {
+          result.push_back(getDynamicSizes()[ctr++]);
+        } else {
+          result.push_back(b.getIndexAttr(getType().getShape()[i]));
+        }
+      }
+      return result;
+    }
   }];
 
   let assemblyFormat = [{
     `(`$dynamicSizes`)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
 
   let hasCanonicalizer = 1;
   let hasVerifier = 1;

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterfaceImpl.cpp

 //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterfaceImpl.h"
 
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using presburger::IntegerPolyhedron;
 
 namespace mlir {
 namespace tensor {
SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
         applyOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
     AffineExpr bound =
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
     cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, bound);
   };
 };
 } // namespace
 
+namespace memref {
+namespace {
+
+template <typename OpTy>
+struct AllocOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<AllocOpInterface<OpTy>,
+                                                   OpTy> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto allocOp = cast<OpTy>(op);
+    assert(value == allocOp.getResult() && "invalid value");
+
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                  cstr.getExpr(allocOp.getMixedSizes()[dim]));
+  }
+};
+
+struct CastOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<CastOpInterface, CastOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto castOp = cast<CastOp>(op);
+    assert(value == castOp.getResult() && "invalid value");
+
+    if (castOp.getResult().getType().isa<MemRefType>() &&
+        castOp.getSource().getType().isa<MemRefType>()) {
+      cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                    cstr.getExpr(castOp.getSource(), dim));
+    }
+  }
+};
+
+struct DimOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<DimOpInterface, DimOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto dimOp = cast<DimOp>(op);
+    assert(value == dimOp.getResult() && "invalid value");
+
+    auto constIndex = dimOp.getConstantIndex();
+    if (!constIndex.has_value())
+      return;
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(dimOp.getSource(), *constIndex));
+  }
+};
+
+struct GetGlobalOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<GetGlobalOpInterface,
+                                                   GetGlobalOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto getGlobalOp = cast<GetGlobalOp>(op);
+    assert(value == getGlobalOp.getResult() && "invalid value");
+
+    auto type = getGlobalOp.getType();
+    assert(!type.isDynamicDim(dim) && "expected static dim");
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                  cstr.getExpr(type.getDimSize(dim)));
+  }
+};
+
+struct RankOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<RankOpInterface, RankOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto rankOp = cast<RankOp>(op);
+    assert(value == rankOp.getResult() && "invalid value");
+
+    auto memrefType = rankOp.getMemref().getType().dyn_cast<MemRefType>();
+    if (!memrefType)
+      return;
+    cstr.addBound(IntegerPolyhedron::BoundType::EQ, value,
+                  cstr.getExpr(memrefType.getRank()));
+  }
+};
+
+struct SubViewOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<SubViewOpInterface,
+                                                   SubViewOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto subViewOp = cast<SubViewOp>(op);
+    assert(value == subViewOp.getResult() && "invalid value");
+
+    llvm::SmallBitVector dropped = subViewOp.getDroppedDims();
+    int64_t ctr = -1;
+    for (int64_t i = 0, e = subViewOp.getMixedSizes().size(); i < e; ++i) {
+      // Skip over rank-reduced dimensions.
+      if (!dropped.test(i))
+        ++ctr;
+      if (ctr == dim) {
+        cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim,
+                      subViewOp.getMixedSizes()[i]);
+        return;
+      }
+    }
+    llvm_unreachable("could not find non-rank-reduced dim");
+  }
+};
+
+} // namespace
+} // namespace memref
 } // namespace mlir
 
 void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
     tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
     tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
     tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
     tensor::ExtractSliceOp::attachInterface<tensor::ExtractSliceOpInterface>(
         *ctx);
     tensor::InsertOp::attachInterface<
         DstValueBoundsOpInterfaceExternalModel<tensor::InsertOp>>(*ctx);
     tensor::InsertSliceOp::attachInterface<
         DstValueBoundsOpInterfaceExternalModel<tensor::InsertSliceOp>>(*ctx);
     tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
     tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
   });
 
   registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
     AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
   });
+
+  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+    memref::AllocOp::attachInterface<memref::AllocOpInterface<memref::AllocOp>>(
+        *ctx);
+    memref::AllocaOp::attachInterface<
+        memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
+    memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
+    memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
+    memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(*ctx);
+    memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
+    memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
+  });
 }

mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir

 //       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]], %[[b]]]
 //       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]], %[[b]]]
 //       CHECL:   return %[[apply]]
 func.func @affine_apply(%a: index, %b: index) -> index {
   %0 = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%a, %b]
   %1 = "test.reify_bound"(%0) : (index) -> (index)
   return %1 : index
 }
+
+// -----
+
+// CHECK-LABEL: func @memref_alloc(
+//  CHECK-SAME:     %[[sz:.*]]: index
+//       CHECK:   %[[c6:.*]] = arith.constant 6 : index
+//       CHECK:   return %[[c6]], %[[sz]]
+func.func @memref_alloc(%sz: index) -> (index, index) {
+  %0 = memref.alloc(%sz) : memref<6x?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (memref<6x?xf32>) -> (index)
+  %2 = "test.reify_bound"(%0) {dim = 1} : (memref<6x?xf32>) -> (index)
+  return %1, %2 : index, index
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_alloca(
+//  CHECK-SAME:     %[[sz:.*]]: index
+//       CHECK:   %[[c6:.*]] = arith.constant 6 : index
+//       CHECK:   return %[[c6]], %[[sz]]
+func.func @memref_alloca(%sz: index) -> (index, index) {
+  %0 = memref.alloca(%sz) : memref<6x?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (memref<6x?xf32>) -> (index)
+  %2 = "test.reify_bound"(%0) {dim = 1} : (memref<6x?xf32>) -> (index)
+  return %1, %2 : index, index
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_cast(
+//       CHECK:   %[[c10:.*]] = arith.constant 10 : index
+//       CHECK:   return %[[c10]]
+func.func @memref_cast(%m: memref<10xf32>) -> index {
+  %0 = memref.cast %m : memref<10xf32> to memref<?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (memref<?xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_dim(
+//  CHECK-SAME:     %[[m:.*]]: memref<?xf32>
+//       CHECK:   %[[dim:.*]] = memref.dim %[[m]]
+//       CHECK:   %[[dim:.*]] = memref.dim %[[m]]
+//       CHECK:   return %[[dim]]
+func.func @memref_dim(%m: memref<?xf32>) -> index {
+  %c0 = arith.constant 0 : index
+  %0 = memref.dim %m, %c0 : memref<?xf32>
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_get_global(
+//       CHECK:   %[[c4:.*]] = arith.constant 4 : index
+//       CHECK:   return %[[c4]]
+memref.global "private" @gv0 : memref<4xf32> = dense<[0.0, 1.0, 2.0, 3.0]>
+func.func @memref_get_global() -> index {
+  %0 = memref.get_global @gv0 : memref<4xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (memref<4xf32>) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_rank(
+//  CHECK-SAME:     %[[t:.*]]: memref<5xf32>
+//       CHECK:   %[[c1:.*]] = arith.constant 1 : index
+//       CHECK:   return %[[c1]]
+func.func @memref_rank(%m: memref<5xf32>) -> index {
+  %0 = memref.rank %m : memref<5xf32>
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_subview(
+//  CHECK-SAME:     %[[m:.*]]: memref<?xf32>, %[[sz:.*]]: index
+//       CHECK:   return %[[sz]]
+func.func @memref_subview(%m: memref<?xf32>, %sz: index) -> index {
+  %0 = memref.subview %m[2][%sz][1] : memref<?xf32> to memref<?xf32, strided<[1], offset: 2>>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (memref<?xf32, strided<[1], offset: 2>>) -> (index)
+  return %1 : index
+}