diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h --- a/mlir/include/mlir/Conversion/Passes.h +++ b/mlir/include/mlir/Conversion/Passes.h @@ -20,6 +20,7 @@ #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" diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td --- a/mlir/include/mlir/Conversion/Passes.td +++ b/mlir/include/mlir/Conversion/Passes.td @@ -209,6 +209,38 @@ 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 //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h b/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h @@ -0,0 +1,47 @@ +//===- 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 +#include + +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 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 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 createOpCleanerPass(); + +} // namespace mlir + +#endif // MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_ diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt --- a/mlir/lib/Conversion/CMakeLists.txt +++ b/mlir/lib/Conversion/CMakeLists.txt @@ -12,6 +12,7 @@ add_subdirectory(SCFToGPU) add_subdirectory(SCFToSPIRV) add_subdirectory(SCFToStandard) +add_subdirectory(SCFToAffine) add_subdirectory(ShapeToStandard) add_subdirectory(SPIRVToLLVM) add_subdirectory(StandardToLLVM) diff --git a/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt @@ -0,0 +1,17 @@ +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 + ) diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp @@ -0,0 +1,363 @@ +//===- 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(op)) { + if (forOp.getNumResults() > + 0) // JC: affine.for with results seems not supported? + return true; + else if (auto cst = forOp.step().getDefiningOp()) { + // 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 { + void runOnOperation() override; +}; + +// Raise scf.for to affine.for & replace scf.yield with affine.yield +struct SCFForRaising : public OpRewritePattern { + using OpRewritePattern::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(step.getDefiningOp()) + .getValue() + .cast() + .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(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(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(ctx); +} + +void SCFForRaisingPass::runOnOperation() { + OwningRewritePatternList patterns; + SCFForRaisingPatterns(patterns, &getContext()); + SCFForRaisingTarget target(getContext()); + target.addLegalDialect(); + target.addDynamicallyLegalOp(); + target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); + if (failed(applyPartialConversion(getOperation(), target, patterns))) + signalPassFailure(); +} + +std::unique_ptr mlir::createRaiseSCFForPass() { + return std::make_unique(); +} + +// 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(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(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(op)) { + for (auto index : loadOp.getIndices()) { + if (!isAffineType(index)) + return true; + } + return false; + } else if (auto storeOp = dyn_cast(op)) { + for (auto index : storeOp.getIndices()) { + if (!isAffineType(index)) + return true; + } + return false; + } else + return true; + } +}; + +// Rewriting Pass +struct LoadStoreRaisingPass + : public LoadStoreRaisingBase { + void runOnOperation() override; +}; + +// Raise std.load to affine.load +struct LoadRaising : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(LoadOp loadOp, + PatternRewriter &rewriter) const override; +}; + +// Raise std.store to affine.store +struct StoreRaising : public OpRewritePattern { + using OpRewritePattern::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 *dims, + std::vector *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(op)) + return getAffineExpr(op->getOperand(0), rewriter, dims, symbols) + + getAffineExpr(op->getOperand(1), rewriter, dims, symbols); + else if (isa(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 indices, dims, symbols; + std::vector exprs; + + for (auto index : loadOp.getIndices()) { + indices.push_back(index); + exprs.push_back(getAffineExpr(index, rewriter, &dims, &symbols)); + } + ArrayRef results(exprs); + AffineMap affineMap = AffineMap::get(dims.size(), symbols.size(), results, + rewriter.getContext()); + dims.insert(dims.end(), symbols.begin(), symbols.end()); + rewriter.replaceOpWithNewOp(loadOp, loadOp.getMemRef(), + affineMap, dims); + return success(); +} + +// Raise std.store to affine.store +LogicalResult StoreRaising::matchAndRewrite(StoreOp storeOp, + PatternRewriter &rewriter) const { + + std::vector indices, dims, symbols; + std::vector exprs; + + for (auto index : storeOp.getIndices()) { + indices.push_back(index); + exprs.push_back(getAffineExpr(index, rewriter, &dims, &symbols)); + } + ArrayRef results(exprs); + AffineMap affineMap = AffineMap::get(dims.size(), symbols.size(), results, + rewriter.getContext()); + dims.insert(dims.end(), symbols.begin(), symbols.end()); + rewriter.replaceOpWithNewOp( + storeOp, storeOp.getValueToStore(), storeOp.getMemRef(), affineMap, dims); + return success(); +} + +void mlir::LoadStoreRaisingPatterns(OwningRewritePatternList &patterns, + MLIRContext *ctx) { + patterns.insert(ctx); +} + +void LoadStoreRaisingPass::runOnOperation() { + OwningRewritePatternList patterns; + LoadStoreRaisingPatterns(patterns, &getContext()); + LoadStoreRaisingTarget target(getContext()); + target.addLegalDialect(); + target.addDynamicallyLegalOp(); + target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); + if (failed(applyPartialConversion(getOperation(), target, patterns))) + signalPassFailure(); +} + +std::unique_ptr mlir::createRaiseLoadStorePass() { + return std::make_unique(); +} + +// 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(op) || isa(op))) { + for (const mlir::OpResult &result : op->getResults()) { + if (!result.use_empty()) + return true; + } + return false; + } else + return true; + } +}; + +struct OpCleanerPass : public OpCleanerBase { + 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 { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(AddIOp addIOp, + PatternRewriter &rewriter) const override { + rewriter.eraseOp(addIOp); + return success(); + }; +}; + +struct MulIOpCleaning : public OpRewritePattern { + using OpRewritePattern::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(ctx); +} + +void OpCleanerPass::runOnOperation() { + OwningRewritePatternList patterns; + OpCleanerPatterns(patterns, &getContext()); + OpCleanerTarget target(getContext()); + target.addDynamicallyLegalDialect(); + target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); + if (failed(applyPartialConversion(getOperation(), target, patterns))) + signalPassFailure(); +} + +/// Create a OpCleaner pass. +std::unique_ptr mlir::createOpCleanerPass() { + return std::make_unique(); +} diff --git a/mlir/test/Conversion/SCFToAffine/raise-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/raise-to-affine.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/SCFToAffine/raise-to-affine.mlir @@ -0,0 +1,307 @@ +// 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 +}