[mlir][linalg] Add ValueBoundsOpInterface impls for Linalg structured ops

[springerm]

Summary

Also add a few more complex test cases.

Depends On: D145804

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[dcaballe]

mlir/test/Dialect/Affine/value-bounds-reification.mlir

78

Awesome!

mlir/include/mlir/Dialect/Linalg/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_LINALG_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_LINALG_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace linalg {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace linalg
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LINALG_IR_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 #include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.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"
inline void registerAllDialects(DialectRegistry &registry) {
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
+  linalg::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerValueBoundsOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerValueBoundsOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);

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

 add_mlir_dialect_library(MLIRLinalgDialect
   LinalgInterfaces.cpp
   LinalgOps.cpp
   LinalgDialect.cpp
+  ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgInterfacesIncGen
   MLIRLinalgOpsAttributesIncGen
   MLIRLinalgOpsEnumsIncGen
add_mlir_dialect_library(MLIRLinalgDialect
   MLIRParser
   MLIRSideEffectInterfaces
   MLIRSparseTensorDialect
   MLIRSCFDialect
   MLIRMathDialect
   MLIRMemRefDialect
   MLIRTensorDialect
   MLIRTilingInterface
+  MLIRValueBoundsOpInterface
   MLIRViewLikeInterface
   )

mlir/lib/Dialect/Linalg/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/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using presburger::BoundType;
+
+namespace mlir {
+namespace linalg {
+namespace {
+
+/// Helper structure that iterates over all LinalgOps in `OpTys` and registers
+/// the `ValueBoundsOpInterface` with each of them.
+template <typename... Ops>
+struct LinalgValueBoundsOpInterfaceHelper {
+  static void registerOpInterface(MLIRContext *ctx) {
+    (Ops::template attachInterface<DstValueBoundsOpInterfaceExternalModel<Ops>>(
+         *ctx),
+     ...);
+  }
+};
+
+} // namespace
+} // namespace linalg
+} // namespace mlir
+
+void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+    // Register all Linalg structured ops.
+    LinalgValueBoundsOpInterfaceHelper<
+#define GET_OP_LIST
+#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
+        >::registerOpInterface(ctx);
+  });
+}

mlir/test/Dialect/Affine/value-bounds-reification.mlir

func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {
   %3 = tensor.dim %2, %c2 : tensor<?x?x?xf32>
 
   %4 = "test.reify_bound"(%2) {dim = 0} : (tensor<?x?x?xf32>) -> (index)
   %5 = "test.reify_bound"(%2) {dim = 1} : (tensor<?x?x?xf32>) -> (index)
   %6 = "test.reify_bound"(%3) : (index) -> (index)
 
   return %4, %5, %6 : index, index, index
 }
+
+// -----
+
+// CHECK-LABEL: func @reify_slice_bound(
+//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK:   "test.some_use"(%[[c5]])
+func.func @reify_slice_bound(%t: tensor<?x?xi32>, %idx: index, %ub: index, %f: f32) {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  scf.for %iv = %c0 to %ub step %c4 {
+    %sz = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 4)>(%iv)[%ub]
+    %slice = tensor.extract_slice %t[%idx, %iv] [1, %sz] [1, 1] : tensor<?x?xi32> to tensor<1x?xi32>
+    %filled = linalg.fill ins(%f : f32) outs(%slice : tensor<1x?xi32>) -> tensor<1x?xi32>
+    %bound = "test.reify_bound"(%filled) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
+    "test.some_use"(%bound) : (index) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1 + 1)>
+// CHECK-LABEL: func @scf_for(
+//  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index
+//       CHECK:   %[[bound:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]
+//       CHECK:   "test.some_use"(%[[bound]])
+func.func @scf_for(%lb: index, %ub: index, %step: index) {
+  scf.for %iv = %lb to %ub step %step {
+    %0 = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv)[%ub]
+    %bound = "test.reify_bound"(%0) {type = "UB"} : (index) -> (index)
+    "test.some_use"(%bound) : (index) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @reify_slice_bound2(
+func.func @reify_slice_bound2(%lb0: index, %ub0: index, %step0: index,
+                              %ub2: index, %t1: tensor<1x?xi8>,
+                              %t2: tensor<?x?xi8>, %t3: tensor<1x?xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c32 = arith.constant 32 : index
+  scf.for %iv0 = %lb0 to %ub0 step %step0 {
+    // CHECK: %[[c129:.*]] = arith.constant 129 : index
+    // CHECK: "test.some_use"(%[[c129]])
+    %ub1 = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 128)>(%iv0)[%ub0]
+    %ub1_ub = "test.reify_bound"(%ub1) {type = "UB"} : (index) -> (index)
+    "test.some_use"(%ub1_ub) : (index) -> ()
+
+    // CHECK: %[[c129:.*]] = arith.constant 129 : index
+    // CHECK: "test.some_use"(%[[c129]])
+    %lb1 = affine.apply affine_map<()[s0] -> ((s0 floordiv 32) * 32)>()[%ub1]
+    %lb1_ub = "test.reify_bound"(%lb1) {type = "UB"} : (index) -> (index)
+    "test.some_use"(%lb1_ub) : (index) -> ()
+
+    scf.for %iv1 = %lb1 to %ub1 step %c32 {
+      // CHECK: %[[c32:.*]] = arith.constant 32 : index
+      // CHECK: "test.some_use"(%[[c32]])
+      %sz = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv1)[%ub1]
+      %sz_ub = "test.reify_bound"(%sz) {type = "UB"} : (index) -> (index)
+      "test.some_use"(%sz_ub) : (index) -> ()
+
+      scf.for %iv2 = %c0 to %ub2 step %c1 {
+        %slice1 = tensor.extract_slice %t1[0, %iv2] [1, 1] [1, 1] : tensor<1x?xi8> to tensor<1x1xi8>
+        %slice2 = tensor.extract_slice %t2[%iv2, 0] [1, %sz] [1, 1] : tensor<?x?xi8> to tensor<1x?xi8>
+        %slice3 = tensor.extract_slice %t3[0, 0] [1, %sz] [1, 1] : tensor<1x?xi32> to tensor<1x?xi32>
+        %matmul = linalg.matmul ins(%slice1, %slice2 : tensor<1x1xi8>, tensor<1x?xi8>) outs(%slice3 : tensor<1x?xi32>) -> tensor<1x?xi32>
+
+        // CHECK: %[[c32:.*]] = arith.constant 32 : index
+        // CHECK: "test.some_use"(%[[c32]])
+        %matmul_ub = "test.reify_bound"(%matmul) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
+        "test.some_use"(%matmul_ub) : (index) -> ()
+      }
+    }
+  }
+  return
+}

mlir/test/Dialect/Linalg/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 @linalg_fill(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
+//       CHECK:   return %[[dim]]
+func.func @linalg_fill(%t: tensor<?xf32>, %f: f32) -> index {
+  %0 = linalg.fill ins(%f : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}

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

deps = [
         ":MemRefDialect",
         ":Parser",
         ":SCFDialect",
         ":SideEffectInterfaces",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         ":TilingInterface",
+        ":ValueBoundsOpInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "LinalgTransformOps",
     srcs = [