[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

30

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

[springerm]

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

30

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/include/mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h

+//===- ValueBoundsOpInterfaceImpl.h - 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MEMREF_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_MEMREF_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace memref {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace memref
+} // namespace mlir
+
+#endif // MLIR_DIALECT_MEMREF_IR_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 #include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
inline void registerAllDialects(DialectRegistry &registry) {
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
+  memref::registerValueBoundsOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   tensor::registerValueBoundsOpInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/MemRef/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefDialect
   MemRefDialect.cpp
   MemRefOps.cpp
+  ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
 
   DEPENDS
   MLIRMemRefOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRArithUtils
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
+  MLIRValueBoundsOpInterface
   MLIRViewLikeInterface
 )

mlir/lib/Dialect/MemRef/IR/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/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using presburger::BoundType;
+
+namespace mlir {
+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(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(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(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(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(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(BoundType::EQ, value, dim, subViewOp.getMixedSizes()[i]);
+        return;
+      }
+    }
+    llvm_unreachable("could not find non-rank-reduced dim");
+  }
+};
+
+} // namespace
+} // namespace memref
+} // namespace mlir
+
+void mlir::memref::registerValueBoundsOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  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/MemRef/value-bounds-op-interface-impl.mlir

+// RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \
+// RUN:     -split-input-file | FileCheck %s
+
+// 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
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

srcs = glob(
         [
             "lib/Dialect/MemRef/IR/*.cpp",
             "lib/Dialect/MemRef/IR/*.h",
             "lib/Dialect/MemRef/Utils/*.cpp",
         ],
     ),
     hdrs = [
         "include/mlir/Dialect/MemRef/IR/MemRef.h",
+        "include/mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h",
         "include/mlir/Dialect/MemRef/Utils/MemRefUtils.h",
     ],
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":ArithUtils",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefBaseIncGen",
         ":MemRefOpsIncGen",
         ":ShapedOpInterfaces",
+        ":ValueBoundsOpInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
         "//llvm:TargetParser",
     ],
 )
 
 gentbl_cc_library(
     name = "MemRefPassIncGen",