[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

2362

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

2368

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

[bondhugula]

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

2362

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.

2368

Sure.

[bondhugula]

Add helper.

[bondhugula]

Adjust accessor and comments.

[bondhugula]

Trim includes.

[dcaballe]

mlir/test/Dialect/Affine/canonicalize.mlir

990–991

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

[bondhugula]

Drop unnecessary -DAGs.

mlir/test/Dialect/Affine/canonicalize.mlir

990–991

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/lib/Dialect/Affine/IR/AffineOps.cpp

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();
+
+  // Fold load from a global constant memref if load indices are all constant.
+  if (getAffineMap().isConstant()) {
+    if (auto getGlobalOp = memref().getDefiningOp<memref::GetGlobalOp>()) {
+      // Get to the memref.global defining the symbol.
+      if (auto *symbolTableOp =
+              getGlobalOp->getParentWithTrait<OpTrait::SymbolTable>()) {
+        if (auto global =
+                dyn_cast_or_null<memref::GlobalOp>(SymbolTable::lookupSymbolIn(
+                    symbolTableOp, getGlobalOp.nameAttr()))) {
+
+          // Return the constant for a memref.global if constant or null
+          // otherwise.
+          auto getGlobalConstantValue =
+              [&](memref::GlobalOp globalOp) -> DenseElementsAttr {
+            auto initVal = globalOp.initial_value();
+            if (!globalOp.constant() || !initVal.hasValue())
+              return {};
+            return initVal.getValue().dyn_cast<DenseElementsAttr>();
+          };
+
+          // Check for constant global memref and access value at known indices.
+          if (auto cstAttr = getGlobalConstantValue(global)) {
+            auto indices = llvm::to_vector<4>(
+                llvm::map_range(getAffineMap().getConstantResults(),
+                                [](int64_t v) -> uint64_t { return v; }));
+            return cstAttr.getValues<Attribute>()[indices];
+          }
+        }
+      }
+    }
+  }
   return OpFoldResult();
 }
 
 //===----------------------------------------------------------------------===//
 // AffineStoreOp
 //===----------------------------------------------------------------------===//
 
 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,

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>
+    // CHECK-DAG: %[[C0:.*]] = arith.constant 2.500000e-01 : f32
+    // CHECK-DAG: %[[C1:.*]] = arith.constant 6.250000e-02 : f32
+    // CHECK-NEXT: return
+    // CHECK-DAG-SAME: %[[C0]]
+    // CHECK-DAG-SAME: %[[C1]]
+    return %v0, %v1 : f32, f32
+  }
+}