added support raising scf to affine dialect

[JianyiCheng]

Summary

• raise scf.for and scf.yield to affine.for and affine.yield ops
• raise std.load and std.store to affine.load and affine.store ops

[arnab-oss]

ping @ JianyiCheng, can you please let me know if this revision has been abandoned? I actually require a similar kind of pass.

[JianyiCheng]

In D90333#3225399, @arnab-oss wrote:

ping @ JianyiCheng, can you please let me know if this revision has been abandoned? I actually require a similar kind of pass.

Hi, this patch has been merged into Polygiest now: https://polygeist.mit.edu, and I should have closed this... sorry about that!

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
+#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"

mlir/include/mlir/Conversion/Passes.td

 def SCFToStandard : Pass<"convert-scf-to-std"> {
   let summary = "Convert SCF dialect to Standard dialect, replacing structured"
                 " control flow with a CFG";
   let constructor = "mlir::createLowerToCFGPass()";
   let dependentDialects = ["StandardOpsDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// SCFForRaising
+//===----------------------------------------------------------------------===//
+
+def SCFForRaising : Pass<"scffor-raising"> {
+  let summary = "Convert scf.for and scf.yield ops to affine.for and "
+                "affine.yield ops";
+  let constructor = "mlir::createRaiseSCFForPass()";
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect", "StandardOpsDialect"];
+}
+
+//===----------------------------------------------------------------------===//
+// LoadStoreRaising
+//===----------------------------------------------------------------------===//
+
+def LoadStoreRaising : Pass<"load-store-raising"> {
+  let summary = "Convert std.load and std.store ops to  affine.load and "
+                "affine.store ops";
+  let constructor = "mlir::createRaiseLoadStorePass()";
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect", "StandardOpsDialect"];
+}
+
+//===----------------------------------------------------------------------===//
+// OpCleaner
+//===----------------------------------------------------------------------===//
+
+def OpCleaner : Pass<"clean"> {
+  let summary = "remove unused ops";
+  let constructor = "mlir::createOpCleanerPass()";
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect", "StandardOpsDialect"];
+}
+
+//===----------------------------------------------------------------------===//
 // SCFToGPU
 //===----------------------------------------------------------------------===//
 
 def ConvertAffineForToGPU : FunctionPass<"convert-affine-for-to-gpu"> {
   let summary = "Convert top-level AffineFor Ops to GPU kernels";
   let constructor = "mlir::createAffineForToGPUPass()";
   let dependentDialects = ["gpu::GPUDialect"];
   let options = [

mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h

+//===- ConvertSCFToAffine.h - Pass entrypoint -----------------*- C++ -*-===//
+//
+// 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_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
+#define MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
+
+#include <memory>
+#include <vector>
+
+namespace mlir {
+struct LogicalResult;
+class MLIRContext;
+class Pass;
+class RewritePattern;
+
+// Owning list of rewriting patterns.
+class OwningRewritePatternList;
+
+/// Collect a set of patterns to raise scf.for and scf.yield
+/// to affine.for and affine.yield
+void SCFForRaisingPatterns(OwningRewritePatternList &patterns,
+                           MLIRContext *ctx);
+/// Creates a pass to convert scf.for and scf.yield ops
+/// to affine.for and affine.yield
+std::unique_ptr<Pass> createRaiseSCFForPass();
+
+/// Collect a set of patterns to raise std.load and std.store
+/// to affine.load and affine.store
+void LoadStoreRaisingPatterns(OwningRewritePatternList &patterns,
+                              MLIRContext *ctx);
+/// Creates a pass to convert std.load and std.store to
+/// affine.load and affine.store.
+std::unique_ptr<Pass> createRaiseLoadStorePass();
+
+/// Collect a set of patterns to remove unused ops
+void OpCleanerPatterns(OwningRewritePatternList &patterns, MLIRContext *ctx);
+/// Creates a pass to remove unused ops
+std::unique_ptr<Pass> createOpCleanerPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AVX512ToLLVM)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(OpenMPToLLVM)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToSPIRV)
 add_subdirectory(SCFToStandard)
+add_subdirectory(SCFToAffine)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)

mlir/lib/Conversion/SCFToAffine/CMakeLists.txt

+add_mlir_conversion_library(MLIRSCFToAffine
+  SCFToAffine.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/SCFToAffine
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRSCF
+  MLIRAffineOps
+  MLIRTransforms
+  )

mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp

+//===- SCFToAffine.cpp - Raising SCF to Affine conversion
+//------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to convert scf.for, scf.yield, std.load
+// and std.store ops into affine.for, affine.yield, affine.load and
+// affine.store ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
+#include "../PassDetail.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/Utils.h"
+
+using namespace mlir;
+using namespace mlir::scf;
+
+// Raise scf.for and scf.yield to affine.for and affine.yield
+// to enable analysis on dimensions and symbols
+namespace {
+// Conversion Target
+class SCFForRaisingTarget : public ConversionTarget {
+public:
+  explicit SCFForRaisingTarget(MLIRContext &context)
+      : ConversionTarget(context) {}
+
+  bool isDynamicallyLegal(Operation *op) const override {
+    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+      if (forOp.getNumResults() >
+          0) // JC: affine.for with results seems not supported?
+        return true;
+      else if (auto cst = forOp.step().getDefiningOp<ConstantIndexOp>()) {
+        // step > 0 already verified in SCF dialect
+        auto lb = forOp.lowerBound();
+        auto ub = forOp.upperBound();
+        return !(isValidSymbol(lb) && isValidSymbol(ub));
+      } else // only constant step supported in affine.for
+        return true;
+    } else
+      return true;
+  }
+};
+
+// Rewriting Pass
+struct SCFForRaisingPass : public SCFForRaisingBase<SCFForRaisingPass> {
+  void runOnOperation() override;
+};
+
+// Raise scf.for to affine.for & replace scf.yield with affine.yield
+struct SCFForRaising : public OpRewritePattern<ForOp> {
+  using OpRewritePattern<ForOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForOp forOp,
+                                PatternRewriter &rewriter) const override;
+};
+} // namespace
+
+// Raise scf.for to affine.for & replace scf.yield with affine.yield
+// Conversion:
+// Raise scf.for to affine.for:
+// scf.for step -> affine.for step
+// scf.for lb -> affine.for lb
+// scf.for ub -> affine.for ub
+// scf.for iter_ags -> affine.for iter_args
+// scf.yield -> affine.yield
+// scf.for body -> affine.for body
+// create a map for both lb, ub: #map1 = affine_map<()[s0] -> (s0)>
+LogicalResult SCFForRaising::matchAndRewrite(ForOp forOp,
+                                             PatternRewriter &rewriter) const {
+  Location loc = forOp.getLoc();
+  // map scf.for operands to affine.for
+  auto step = forOp.step();
+  auto lb = forOp.lowerBound();
+  auto ub = forOp.upperBound();
+  // auto iterArgs = forOp.getRegionIterArgs();
+  // affine.for only accepts a step as a literal
+  int stepNum = dyn_cast<ConstantOp>(step.getDefiningOp())
+                    .getValue()
+                    .cast<IntegerAttr>()
+                    .getInt();
+  // the loop bounds are both valid symbols - direct map #map <()[s0] -> (s0)>
+  AffineMap directSymbolMap =
+      AffineMap::get(0, 1, getAffineSymbolExpr(0, rewriter.getContext()));
+  auto f = rewriter.create<mlir::AffineForOp>(loc, lb, directSymbolMap, ub,
+                                              directSymbolMap, stepNum);
+  rewriter.eraseBlock(f.getBody());
+  Operation *loopTerminator = forOp.region().back().getTerminator();
+  ValueRange terminatorOperands = loopTerminator->getOperands();
+  rewriter.setInsertionPointToEnd(&forOp.region().back());
+  rewriter.create<AffineYieldOp>(loc, terminatorOperands);
+  rewriter.inlineRegionBefore(forOp.region(), f.region(), f.region().end());
+  rewriter.eraseOp(loopTerminator);
+  rewriter.eraseOp(forOp);
+  return success();
+}
+
+void mlir::SCFForRaisingPatterns(OwningRewritePatternList &patterns,
+                                 MLIRContext *ctx) {
+  patterns.insert<SCFForRaising>(ctx);
+}
+
+void SCFForRaisingPass::runOnOperation() {
+  OwningRewritePatternList patterns;
+  SCFForRaisingPatterns(patterns, &getContext());
+  SCFForRaisingTarget target(getContext());
+  target.addLegalDialect<SCFDialect, AffineDialect>();
+  target.addDynamicallyLegalOp<scf::ForOp>();
+  target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+  if (failed(applyPartialConversion(getOperation(), target, patterns)))
+    signalPassFailure();
+}
+
+std::unique_ptr<Pass> mlir::createRaiseSCFForPass() {
+  return std::make_unique<SCFForRaisingPass>();
+}
+
+// Raise std.load and std.store to affine.load and affine.store
+// after the affineScope is constructed by raising the scf.for loops
+namespace {
+// Conversion Target
+class LoadStoreRaisingTarget : public ConversionTarget {
+public:
+  explicit LoadStoreRaisingTarget(MLIRContext &context)
+      : ConversionTarget(context) {}
+
+  // return whether the index is in affine expression
+  bool isAffineType(Value value) const {
+    if (isValidDim(value) || isValidSymbol(value))
+      return true;
+    else if (auto op = value.getDefiningOp()) {
+      if (isa<MulIOp>(op)) { // cannot recognise the pattern s0*(d1+s2) so far
+        if (isValidSymbol(op->getOperand(0)))
+          return isValidDim(op->getOperand(1));
+        else if (isValidSymbol(op->getOperand(1)))
+          return isValidDim(op->getOperand(0));
+        else
+          return false;
+      } else if (isa<AddIOp>(op)) // affine + affine is affine
+        return (isAffineType(op->getOperand(0)) &&
+                isAffineType(op->getOperand(1)));
+      else
+        return false;
+    } else
+      return false;
+  }
+
+  bool isDynamicallyLegal(Operation *op) const override {
+    if (auto loadOp = dyn_cast<LoadOp>(op)) {
+      for (auto index : loadOp.getIndices()) {
+        if (!isAffineType(index))
+          return true;
+      }
+      return false;
+    } else if (auto storeOp = dyn_cast<StoreOp>(op)) {
+      for (auto index : storeOp.getIndices()) {
+        if (!isAffineType(index))
+          return true;
+      }
+      return false;
+    } else
+      return true;
+  }
+};
+
+// Rewriting Pass
+struct LoadStoreRaisingPass
+    : public LoadStoreRaisingBase<LoadStoreRaisingPass> {
+  void runOnOperation() override;
+};
+
+// Raise std.load to affine.load
+struct LoadRaising : public OpRewritePattern<LoadOp> {
+  using OpRewritePattern<LoadOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(LoadOp loadOp,
+                                PatternRewriter &rewriter) const override;
+};
+
+// Raise std.store to affine.store
+struct StoreRaising : public OpRewritePattern<StoreOp> {
+  using OpRewritePattern<StoreOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(StoreOp storeOp,
+                                PatternRewriter &rewriter) const override;
+};
+} // namespace
+
+// Extract the affine expression from a number of std operations
+AffineExpr getAffineExpr(Value value, PatternRewriter &rewriter,
+                         std::vector<Value> *dims,
+                         std::vector<Value> *symbols) {
+  auto op = value.getDefiningOp();
+  assert(op != NULL || isValidDim(value));
+  if (isValidSymbol(value)) {
+    int symbolIdx;
+    auto symbolIter = std::find(symbols->begin(), symbols->end(), value);
+    if (symbolIter == symbols->end()) {
+      symbolIdx = symbols->size();
+      symbols->push_back(value);
+    } else
+      symbolIdx = std::distance(symbols->begin(), symbolIter);
+    return getAffineSymbolExpr(symbolIdx, rewriter.getContext());
+  } else if (isValidDim(value)) {
+    int dimIdx;
+    auto dimIter = std::find(dims->begin(), dims->end(), value);
+    if (dimIter == dims->end()) {
+      dimIdx = dims->size();
+      dims->push_back(value);
+    } else
+      dimIdx = std::distance(dims->begin(), dimIter);
+    return getAffineDimExpr(dimIdx, rewriter.getContext());
+  } else if (isa<AddIOp>(op))
+    return getAffineExpr(op->getOperand(0), rewriter, dims, symbols) +
+           getAffineExpr(op->getOperand(1), rewriter, dims, symbols);
+  else if (isa<MulIOp>(op))
+    return getAffineExpr(op->getOperand(0), rewriter, dims, symbols) *
+           getAffineExpr(op->getOperand(1), rewriter, dims, symbols);
+  else
+    return NULL;
+}
+
+// Raise std.load to affine.load
+// for each index of the loadOp, extract its affine expression
+// construct a map with the dim and symbol count for the affine.load
+// replace the loadOp with the newly constructed affine.load
+LogicalResult LoadRaising::matchAndRewrite(LoadOp loadOp,
+                                           PatternRewriter &rewriter) const {
+  std::vector<Value> indices, dims, symbols;
+  std::vector<AffineExpr> exprs;
+
+  for (auto index : loadOp.getIndices()) {
+    indices.push_back(index);
+    exprs.push_back(getAffineExpr(index, rewriter, &dims, &symbols));
+  }
+  ArrayRef<AffineExpr> results(exprs);
+  AffineMap affineMap = AffineMap::get(dims.size(), symbols.size(), results,
+                                       rewriter.getContext());
+  dims.insert(dims.end(), symbols.begin(), symbols.end());
+  rewriter.replaceOpWithNewOp<AffineLoadOp>(loadOp, loadOp.getMemRef(),
+                                            affineMap, dims);
+  return success();
+}
+
+// Raise std.store to affine.store
+LogicalResult StoreRaising::matchAndRewrite(StoreOp storeOp,
+                                            PatternRewriter &rewriter) const {
+
+  std::vector<Value> indices, dims, symbols;
+  std::vector<AffineExpr> exprs;
+
+  for (auto index : storeOp.getIndices()) {
+    indices.push_back(index);
+    exprs.push_back(getAffineExpr(index, rewriter, &dims, &symbols));
+  }
+  ArrayRef<AffineExpr> results(exprs);
+  AffineMap affineMap = AffineMap::get(dims.size(), symbols.size(), results,
+                                       rewriter.getContext());
+  dims.insert(dims.end(), symbols.begin(), symbols.end());
+  rewriter.replaceOpWithNewOp<AffineStoreOp>(
+      storeOp, storeOp.getValueToStore(), storeOp.getMemRef(), affineMap, dims);
+  return success();
+}
+
+void mlir::LoadStoreRaisingPatterns(OwningRewritePatternList &patterns,
+                                    MLIRContext *ctx) {
+  patterns.insert<LoadRaising, StoreRaising>(ctx);
+}
+
+void LoadStoreRaisingPass::runOnOperation() {
+  OwningRewritePatternList patterns;
+  LoadStoreRaisingPatterns(patterns, &getContext());
+  LoadStoreRaisingTarget target(getContext());
+  target.addLegalDialect<SCFDialect, AffineDialect>();
+  target.addDynamicallyLegalOp<LoadOp, StoreOp>();
+  target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+  if (failed(applyPartialConversion(getOperation(), target, patterns)))
+    signalPassFailure();
+}
+
+std::unique_ptr<Pass> mlir::createRaiseLoadStorePass() {
+  return std::make_unique<LoadStoreRaisingPass>();
+}
+
+// Removing the operations have results but not used.
+namespace {
+
+class OpCleanerTarget : public ConversionTarget {
+public:
+  explicit OpCleanerTarget(MLIRContext &context) : ConversionTarget(context) {}
+
+  bool isDynamicallyLegal(Operation *op) const override {
+    if (op->getNumResults() > 0 && (isa<AddIOp>(op) || isa<MulIOp>(op))) {
+      for (const mlir::OpResult &result : op->getResults()) {
+        if (!result.use_empty())
+          return true;
+      }
+      return false;
+    } else
+      return true;
+  }
+};
+
+struct OpCleanerPass : public OpCleanerBase<OpCleanerPass> {
+  void runOnOperation() override;
+};
+
+// remove unused addi and muli apps caused by affine expression
+// TODO: make it a generic pass to remove all unused op - seems like
+// the rewrite pattern cannot be generic op (Operation *)
+struct AddIOpCleaning : public OpRewritePattern<AddIOp> {
+  using OpRewritePattern<AddIOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(AddIOp addIOp,
+                                PatternRewriter &rewriter) const override {
+    rewriter.eraseOp(addIOp);
+    return success();
+  };
+};
+
+struct MulIOpCleaning : public OpRewritePattern<MulIOp> {
+  using OpRewritePattern<MulIOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(MulIOp mulIOp,
+                                PatternRewriter &rewriter) const override {
+    rewriter.eraseOp(mulIOp);
+    return success();
+  };
+};
+} // end anonymous namespace
+
+void mlir::OpCleanerPatterns(OwningRewritePatternList &patterns,
+                             MLIRContext *ctx) {
+  patterns.insert<AddIOpCleaning, MulIOpCleaning>(ctx);
+}
+
+void OpCleanerPass::runOnOperation() {
+  OwningRewritePatternList patterns;
+  OpCleanerPatterns(patterns, &getContext());
+  OpCleanerTarget target(getContext());
+  target.addDynamicallyLegalDialect<SCFDialect, AffineDialect,
+                                    StandardOpsDialect>();
+  target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+  if (failed(applyPartialConversion(getOperation(), target, patterns)))
+    signalPassFailure();
+}
+
+/// Create a OpCleaner pass.
+std::unique_ptr<Pass> mlir::createOpCleanerPass() {
+  return std::make_unique<OpCleanerPass>();
+}

mlir/test/Conversion/SCFToAffine/raise-to-affine.mlir

+// RUN: mlir-opt -scffor-raising -load-store-raising -clean %s | FileCheck %s
+  
+// CHECK-LABEL: func @constLoad(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %0 = affine.load %arg0[symbol(%c10)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @constLoad(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %0 = load %arg0[%b] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @constStore(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %c10_i32 = constant 10 : i32
+// CHECK-NEXT:     affine.store %c10_i32, %arg0[symbol(%c10)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @constStore(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %c = constant 10 : i32
+    store %c, %arg0[%b] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @singleDimLoad(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %0 = affine.load %arg0[%arg1] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @singleDimLoad(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %0 = load %arg0[%arg1] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @singleDimStore(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10_i32 = constant 10 : i32
+// CHECK-NEXT:     affine.store %c10_i32, %arg0[%arg1] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @singleDimStore(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %0 = constant 10 : i32
+    store %0, %arg0[%arg1] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @mulILoad(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %0 = affine.load %arg0[%arg1 * symbol(%c10)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @mulILoad(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %a = muli %arg1, %b : index
+    %0 = load %arg0[%a] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @mulIStore(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %c0_i32_0 = constant 0 : i32
+// CHECK-NEXT:     affine.store %c0_i32_0, %arg0[%arg1 * symbol(%c10)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @mulIStore(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %a = muli %arg1, %b : index
+    %0 = constant 0 : i32
+    store %0, %arg0[%a] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @addILoad(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %0 = affine.load %arg0[%arg1 + symbol(%c10)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @addILoad(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %a = addi %arg1, %b : index
+    %0 = load %arg0[%a] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @addIStore(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %c0_i32_0 = constant 0 : i32
+// CHECK-NEXT:     affine.store %c0_i32_0, %arg0[%arg1 + symbol(%c10)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @addIStore(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %a = addi %arg1, %b : index
+    %0 = constant 0 : i32
+    store %0, %arg0[%a] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @nonAffineLoad(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %0 = muli %arg1, %arg1 : index
+// CHECK-NEXT:     %1 = load %arg0[%0] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @nonAffineLoad(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %a = muli %arg1, %arg1 : index
+    %0 = load %arg0[%a] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @loadTest(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %0 = affine.load %arg0[symbol(%c1)] : memref<1000xi32>
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %1 = muli %arg1, %arg1 : index
+// CHECK-NEXT:     %2 = load %arg0[%1] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+func @loadTest(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %8 = load %arg0[%c1] : memref<1000xi32>
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %a = muli %arg1, %arg1 : index
+    %0 = load %arg0[%a] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @storeTest(%arg0: memref<1000xi32>) -> i32 {
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %0 = addi %c1, %c0 : index
+// CHECK-NEXT:   affine.store %c0_i32, %arg0[symbol(%0)] : memref<1000xi32>
+// CHECK-NEXT:   %c1000 = constant 1000 : index
+// CHECK-NEXT:   affine.for %arg1 = %c0 to %c1000 {
+// CHECK-NEXT:     %c10 = constant 10 : index
+// CHECK-NEXT:     %1 = muli %arg1, %c10 : index
+// CHECK-NEXT:     %c20 = constant 20 : index
+// CHECK-NEXT:     %c0_i32_0 = constant 0 : i32
+// CHECK-NEXT:     affine.store %c0_i32_0, %arg0[%arg1 * symbol(%c10) + symbol(%c20)] : memref<1000xi32>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return %c0_i32 : i32
+// CHECK-NEXT: }
+
+
+func @storeTest(%arg0: memref<1000xi32>) -> i32 {
+  %c0_i32 = constant 0 : i32
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %x = addi %c1, %c0 : index
+  store %c0_i32, %arg0[%x] : memref<1000xi32>
+  %c1000 = constant 1000 : index
+  scf.for %arg1 = %c0 to %c1000 step %c1 {
+    %b = constant 10 : index
+    %a = muli %arg1, %b : index
+    %c = constant 20 : index
+    %d = addi %a, %c : index
+    %0 = constant 0 : i32
+    store %0, %arg0[%d] : memref<1000xi32>
+  }
+  return %c0_i32 : i32
+}