[MLIR] Add affine.load fold hook on global constant memrefs

[bondhugula]

Summary

Fold affine.load ops on global constant memrefs when indices are all
constant.

[nicolasvasilache]

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

2404

Plz invert some conditions and use early exit instead of applying 5-deep nesting.

2410

this would look better as a helper on memref::GlobalOp

[bondhugula]

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

2404

I did think of inverting the checks (which would be the normal way); the reason I didn't is that this is a folding hook where the list of conditions may grow (further below this check). Inverting would lead to a reversal when the next folding rule is added.

And you can't invert the inner ones since there isn't a for loop here (there's no early exit).

Fine with inverting though but that would lead to a larger diff on the next change.

2410

Sure.

[bondhugula]

Add helper.

[bondhugula]

Adjust accessor and comments.

[bondhugula]

Trim includes.

[dcaballe]

mlir/test/Dialect/Affine/canonicalize.mlir

1096–1097

Remove DAG? The order in the return should always be the same, right?

[bondhugula]

Drop unnecessary -DAGs.

mlir/test/Dialect/Affine/canonicalize.mlir

1096–1097

Right. The second set of -DAGs aren't needed.

[bondhugula]

Rebase.

[bondhugula]

Rebase. Ping @reviewers.

[ayzhuang]

After the most recent rebase, your change to GlobalOp is lost. Please check.

[bondhugula]

Put back lost change.

[bondhugula]

In D120612#3361343, @ayzhuang wrote:

After the most recent rebase, your change to GlobalOp is lost. Please check.

Thanks very much for spotting that. Restored.

[ayzhuang]

Looks great, thanks!

[bondhugula]

Rebase. Reduce nesting depth of conditionals - NFC.

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

 //===----------------------------------------------------------------------===//
 
 def MemRef_GlobalOp : MemRef_Op<"global", [Symbol]> {
   let summary = "declare or define a global memref variable";
   let description = [{
     The `memref.global` operation declares or defines a named global memref
     variable. The backing memory for the variable is allocated statically and is
     described by the type of the variable (which should be a statically shaped
-    memref type). The operation is a declaration if no `inital_value` is
+    memref type). The operation is a declaration if no `initial_value` is
     specified, else it is a definition. The `initial_value` can either be a unit
     attribute to represent a definition of an uninitialized global variable, or
     an elements attribute to represent the definition of a global variable with
     an initial value. The global variable can also be marked constant using the
     `constant` unit attribute. Writing to such constant global variables is
     undefined.
 
     The global variable can be accessed by using the `memref.get_global` to
let assemblyFormat = [{
        attr-dict
   }];
 
   let extraClassDeclaration = [{
      bool isExternal() { return !initial_value(); }
      bool isUninitialized() {
        return !isExternal() && initial_value().getValue().isa<UnitAttr>();
      }
+     /// Returns the constant initial value if the memref.global is a constant,
+     /// or null otherwise.
+     ElementsAttr getConstantInitValue();
   }];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 

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

 //===- AffineOps.cpp - MLIR Affine Operations -----------------------------===//
 //
 // 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/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
-#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/InliningUtils.h"
-#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "affine-analysis"
 
void AffineLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<SimplifyAffineOp<AffineLoadOp>>(context);
 }
 
 OpFoldResult AffineLoadOp::fold(ArrayRef<Attribute> cstOperands) {
   /// load(memrefcast) -> load
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
-  return OpFoldResult();
+
+  // Fold load from a global constant memref.
+  auto getGlobalOp = memref().getDefiningOp<memref::GetGlobalOp>();
+  if (!getGlobalOp)
+    return {};
+  // Get to the memref.global defining the symbol.
+  auto *symbolTableOp = getGlobalOp->getParentWithTrait<OpTrait::SymbolTable>();
+  if (!symbolTableOp)
+    return {};
+  auto global = dyn_cast_or_null<memref::GlobalOp>(
+      SymbolTable::lookupSymbolIn(symbolTableOp, getGlobalOp.nameAttr()));
+  if (!global)
+    return {};
+  if (auto cstAttr =
+          global.getConstantInitValue().dyn_cast_or_null<DenseElementsAttr>()) {
+    // We can fold only if we know the indices.
+    if (!getAffineMap().isConstant())
+      return {};
+    auto indices = llvm::to_vector<4>(
+        llvm::map_range(getAffineMap().getConstantResults(),
+                        [](int64_t v) -> uint64_t { return v; }));
+    return cstAttr.getValues<Attribute>()[indices];
+  }
+  return {};
 }
 
 //===----------------------------------------------------------------------===//
 // AffineStoreOp
 //===----------------------------------------------------------------------===//
 
 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
                           Value valueToStore, Value memref, AffineMap map,

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

if (!llvm::isPowerOf2_64(alignment))
       return emitError() << "alignment attribute value " << alignment
                          << " is not a power of 2";
   }
 
   // TODO: verify visibility for declarations.
   return success();
 }
 
+ElementsAttr GlobalOp::getConstantInitValue() {
+  auto initVal = initial_value();
+  if (constant() && initVal.hasValue())
+    return initVal.getValue().cast<ElementsAttr>();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // GetGlobalOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 GetGlobalOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   // Verify that the result type is same as the type of the referenced
   // memref.global op.

mlir/test/Dialect/Affine/canonicalize.mlir

func @canonicalize_single_min_max(%i0: index, %i1: index) -> (index, index) {
   %0 = affine.min affine_map<()[s0] -> (s0 + 16)> ()[%i0]
 
   // CHECK-NOT: affine.max
   // CHECK-NEXT: affine.apply #[[$MAP1]]()[%[[I1]]]
   %1 = affine.min affine_map<()[s0] -> (s0 * 4)> ()[%i1]
 
   return %0, %1: index, index
 }
+
+// -----
+
+module {
+  memref.global "private" constant @__constant_1x5x1xf32 : memref<1x5x1xf32> = dense<[[[6.250000e-02], [2.500000e-01], [3.750000e-01], [2.500000e-01], [6.250000e-02]]]>
+  // CHECK-LABEL: func @fold_const_init_global_memref
+  func @fold_const_init_global_memref() -> (f32, f32) {
+    %m = memref.get_global @__constant_1x5x1xf32 : memref<1x5x1xf32>
+    %v0 = affine.load %m[0, 0, 0] : memref<1x5x1xf32>
+    %v1 = affine.load %m[0, 1, 0] : memref<1x5x1xf32>
+    return %v0, %v1 : f32, f32
+    // CHECK-DAG: %[[C0:.*]] = arith.constant 6.250000e-02 : f32
+    // CHECK-DAG: %[[C1:.*]] = arith.constant 2.500000e-01 : f32
+    // CHECK-NEXT: return
+    // CHECK-SAME: %[[C0]]
+    // CHECK-SAME: %[[C1]]
+  }
+}