[mlir][Affine] Add ValueBoundsOpInterface for affine.apply

[springerm]

Summary

Depends On: D145681

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

update API

[dcaballe]

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

37

== -> =? Add a test for this case if it's not covered?

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

8

Use different var names?

[dcaballe]

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

37

nevermind :)

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

8

nevermind

[bondhugula]

@springerm The test suite is failing due to order stability issues related to these changes (or other revisions related to this):

******************** TEST 'MLIR :: Dialect/Tensor/value-bounds-op-interface-impl.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /home/uday/llvm-project-upstream/build/bin/mlir-opt /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir -test-affine-reify-value-bounds -verify-diagnostics      -split-input-file | /home/uday/llvm-project-upstream/build/bin/FileCheck /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir
--
Exit Code: 1

Command Output (stderr):
--
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:107:11: error: CHECK: expected string not found in input
// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
          ^
<stdin>:70:13: note: scanning from here
 return %dim : index
            ^
<stdin>:76:1: note: possible intended match here
#map = affine_map<()[s0, s1] -> (s0 * 2 + s1)>
^
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:113:44: error: undefined variable: $map
// CHECK: %[[bound0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[a]]]

[springerm]

In D145694#4247724, @bondhugula wrote:

@springerm The test suite is failing due to order stability issues related to these changes (or other revisions related to this):

******************** TEST 'MLIR :: Dialect/Tensor/value-bounds-op-interface-impl.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /home/uday/llvm-project-upstream/build/bin/mlir-opt /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir -test-affine-reify-value-bounds -verify-diagnostics      -split-input-file | /home/uday/llvm-project-upstream/build/bin/FileCheck /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir
--
Exit Code: 1

Command Output (stderr):
--
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:107:11: error: CHECK: expected string not found in input
// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
          ^
<stdin>:70:13: note: scanning from here
 return %dim : index
            ^
<stdin>:76:1: note: possible intended match here
#map = affine_map<()[s0, s1] -> (s0 * 2 + s1)>
^
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:113:44: error: undefined variable: $map
// CHECK: %[[bound0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[a]]]

Do you have any extra changes in your build? Does it only happen sometimes or every time you run the test? I cannot reproduce this locally (ran the test 100 times without failure), also no build bot failure email so far.

[bondhugula]

In D145694#4247729, @springerm wrote:

In D145694#4247724, @bondhugula wrote:

@springerm The test suite is failing due to order stability issues related to these changes (or other revisions related to this):

******************** TEST 'MLIR :: Dialect/Tensor/value-bounds-op-interface-impl.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /home/uday/llvm-project-upstream/build/bin/mlir-opt /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir -test-affine-reify-value-bounds -verify-diagnostics      -split-input-file | /home/uday/llvm-project-upstream/build/bin/FileCheck /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir
--
Exit Code: 1

Command Output (stderr):
--
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:107:11: error: CHECK: expected string not found in input
// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
          ^
<stdin>:70:13: note: scanning from here
 return %dim : index
            ^
<stdin>:76:1: note: possible intended match here
#map = affine_map<()[s0, s1] -> (s0 * 2 + s1)>
^
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:113:44: error: undefined variable: $map
// CHECK: %[[bound0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[a]]]

Do you have any extra changes in your build? Does it only happen sometimes or every time you run the test? I cannot reproduce this locally (ran the test 100 times without failure), also no build bot failure email so far.

I have no local changes - it happens on the upstream main branch all the time. We have an unstable container or an order of evaluation dependent code somewhere.

[bondhugula]

In D145694#4248461, @bondhugula wrote:

In D145694#4247729, @springerm wrote:

In D145694#4247724, @bondhugula wrote:

@springerm The test suite is failing due to order stability issues related to these changes (or other revisions related to this):

******************** TEST 'MLIR :: Dialect/Tensor/value-bounds-op-interface-impl.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /home/uday/llvm-project-upstream/build/bin/mlir-opt /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir -test-affine-reify-value-bounds -verify-diagnostics      -split-input-file | /home/uday/llvm-project-upstream/build/bin/FileCheck /home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir
--
Exit Code: 1

Command Output (stderr):
--
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:107:11: error: CHECK: expected string not found in input
// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
          ^
<stdin>:70:13: note: scanning from here
 return %dim : index
            ^
<stdin>:76:1: note: possible intended match here
#map = affine_map<()[s0, s1] -> (s0 * 2 + s1)>
^
/home/uday/llvm-project-upstream/mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir:113:44: error: undefined variable: $map
// CHECK: %[[bound0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[a]]]

Do you have any extra changes in your build? Does it only happen sometimes or every time you run the test? I cannot reproduce this locally (ran the test 100 times without failure), also no build bot failure email so far.

I have no local changes - it happens on the upstream main branch all the time. We have an unstable container or an order of evaluation dependent code somewhere.

This appears to be due to an order of evaluation issue: IR being generated with no guaranteed order. Happens with GCC 8.5.0, GCC 9.4.0, at least. Are you using clang?

[springerm]

In D145694#4248973, @bondhugula wrote:

This appears to be due to an order of evaluation issue: IR being generated with no guaranteed order. Happens with GCC 8.5.0, GCC 9.4.0, at least. Are you using clang?

I found the problem. This happens with gcc only.

The problematic code is:

AffineExpr expr = cstr.getExpr(padOp.getSource(), dim) +
                  cstr.getExpr(padOp.getMixedLowPad()[dim]) +
                  cstr.getExpr(padOp.getMixedHighPad()[dim]);

The evaluation order of the operands in the sum is unspecified in C++. clang evaluates left to right. gcc evaluates right to left. getExpr has a side effect where things a put on a worklist. The order in which things are put on the worklist determines the order in which columns are projected out in the FlatLinearConstraints. Different orders of projection produce different affine_maps during getSliceBounds.

The easy fix here is to split the above statement into 4 statements. But I'm wondering if there is a better to fix this. I think this won't be the last time that we're seeing this issue.

[springerm]

Fix is here: https://github.com/llvm/llvm-project/commit/06d0cd1e81c19cecedbf86356bc7bfb532a1ea80

[bondhugula]

In D145694#4250323, @springerm wrote:

Fix is here: https://github.com/llvm/llvm-project/commit/06d0cd1e81c19cecedbf86356bc7bfb532a1ea80

This is the right fix. We've seen such things in the past. Things generating IR can't be combined in a expression that way. Another common problematic pattern is:
... = b.create<...>(loc, b.create<....>(...), b.create<....>(...));. One can't guarantee here in what order IR will be generated. In general, whenever you have multiple expressions being combined, the individual expressions shouldn't generate IR.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let extraClassDeclaration = [{
 
     /// Returns true if the result of this operation is a symbol in the region
     /// of the closest surrounding op that has the trait AffineScope.
     bool isValidSymbol();
 
     /// Returns true if the result of this operation is a symbol for all its
     /// uses in `region`.
     bool isValidSymbol(Region *region);
+
+    /// Returns all dimension operands.
+    ValueRange getDimOperands() {
+      return OperandRange{getOperands().begin(),
+                          getOperands().begin() + getMap().getNumDims()};
+    }
+
+    /// Returns all symbol operands.
+    ValueRange getSymbolOperands() {
+      return OperandRange{getOperands().begin() + getMap().getNumDims(),
+                          getOperands().end()};
+    }
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 

mlir/include/mlir/Dialect/Affine/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_AFFINE_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_AFFINE_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace affine {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace affine
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_IR_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INITALLDIALECTS_H_
 #define MLIR_INITALLDIALECTS_H_
 
 #include "mlir/Dialect/AMDGPU/AMDGPUDialect.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
inline void registerAllDialects(DialectRegistry &registry) {
   bufferization::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
   memref::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
+  affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);

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

 add_mlir_dialect_library(MLIRAffineDialect
   AffineMemoryOpInterfaces.cpp
   AffineOps.cpp
   AffineValueMap.cpp
+  ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   DEPENDS
   MLIRAffineMemoryOpInterfacesIncGen
   MLIRAffineOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialectUtils
   MLIRIR
   MLIRLoopLikeInterface
   MLIRMemRefDialect
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
+  MLIRValueBoundsOpInterface
   )

mlir/lib/Dialect/Affine/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/Affine/IR/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace {
+
+struct AffineApplyOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<AffineApplyOpInterface,
+                                                   AffineApplyOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto applyOp = cast<AffineApplyOp>(op);
+    assert(value == applyOp.getResult() && "invalid value");
+    assert(applyOp.getAffineMap().getNumResults() == 1 &&
+           "expected single result");
+
+    // Align affine map result with dims/symbols in the constraint set.
+    AffineExpr expr = applyOp.getAffineMap().getResult(0);
+    SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
+        applyOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
+    SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
+        applyOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
+    AffineExpr bound =
+        expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
+    cstr.bound(value) == bound;
+  }
+};
+
+} // namespace
+} // namespace mlir
+
+void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
+    AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
+  });
+}

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

+// RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \
+// RUN:     -split-input-file | FileCheck %s
+
+// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>
+// CHECK-LABEL: func @affine_apply(
+//  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index
+//       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
+}

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

deps = [
         ":ControlFlowInterfaces",
         ":DialectUtils",
         ":IR",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":ShapedOpInterfaces",
         ":SideEffectInterfaces",
         ":Support",
+        ":ValueBoundsOpInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "EmitCDialect",
     srcs = glob([
         "lib/Dialect/EmitC/IR/*.cpp",