TemporaryStorage.h

//===-- Optimizer/Builder/TemporaryStorage.h --------------------*- 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
//
//===----------------------------------------------------------------------===//
//
//  Utility to create an manipulate vector like temporary storages holding
//  Fortran values or descriptors in HLFIR.
//
//  This is useful to deal with array constructors, and temporary storage
//  inside forall and where constructs where it is not known prior to the
//  construct execution how many values will be stored, or where the values
//  at each iteration may have different shapes or type parameters.
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_BUILDER_TEMPORARYSTORAGE_H
#define FORTRAN_OPTIMIZER_BUILDER_TEMPORARYSTORAGE_H
 
#include "flang/Common/idioms.h"
#include "flang/Optimizer/HLFIR/HLFIROps.h"
 
namespace fir {
class FirOpBuilder;
}
 
namespace hlfir {
class Entity;
}
 
namespace fir::factory {

struct Counter {
  Counter(mlir::Location loc, fir::FirOpBuilder &builder,
          mlir::Value initialValue, bool canCountThroughLoops = true);
  mlir::Value getAndIncrementIndex(mlir::Location loc,
                                   fir::FirOpBuilder &builder);
  void reset(mlir::Location loc, fir::FirOpBuilder &builder);
  const bool canCountThroughLoops;
 
private:
  mlir::Value initialValue;
  mlir::Value one;
  mlir::Value index;
};

class HomogeneousScalarStack {
public:
  HomogeneousScalarStack(mlir::Location loc, fir::FirOpBuilder &builder,
                         fir::SequenceType declaredType, mlir::Value extent,
                         llvm::ArrayRef<mlir::Value> lengths,
                         bool allocateOnHeap, bool stackThroughLoops,
                         llvm::StringRef name);
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value);
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);

  hlfir::Entity moveStackAsArrayExpr(mlir::Location loc,
                                     fir::FirOpBuilder &builder);

  bool canBeFetchedAfterPush() const { return false; }
 
private:
  const bool allocateOnHeap;
  Counter counter;
  mlir::Value temp;
};

class ArrayTemp {
public:
  ArrayTemp(mlir::Location loc, fir::FirOpBuilder &builder,
            fir::SequenceType declaredType, llvm::ArrayRef<mlir::Value> extents,
            llvm::ArrayRef<mlir::Value> lowerBounds,
            llvm::ArrayRef<mlir::Value> lengths, bool allocateOnHeap,
            llvm::StringRef name);
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value, mlir::ValueRange indices);
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder) {}
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder,
                    mlir::ValueRange indices);
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
  bool canBeFetchedAfterPush() const { return true; }
 
private:
  const bool allocateOnHeap;
  mlir::Value temp;
  llvm::SmallVector<mlir::Value> typeParams;
};

class SimpleCopy {
public:
  SimpleCopy(mlir::Location loc, fir::FirOpBuilder &builder,
             hlfir::Entity source, llvm::StringRef tempName);
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value) {
    assert(false && "must not be called: value already set");
  }
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder){};
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
    return copy.getBase();
  }
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
  bool canBeFetchedAfterPush() const { return true; }
 
public:
  hlfir::AssociateOp copy;
};

class SSARegister {
public:
  SSARegister(){};
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value) {
    ssaRegister = value;
  }
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder){};
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
    return ssaRegister;
  }
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder) {}
  bool canBeFetchedAfterPush() const { return true; }
 
public:
  mlir::Value ssaRegister;
};

class AnyValueStack {
public:
  AnyValueStack(mlir::Location loc, fir::FirOpBuilder &builder,
                mlir::Type valueStaticType);
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value);
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
  bool canBeFetchedAfterPush() const { return true; }
 
private:
  mlir::Type valueStaticType;
  mlir::Value opaquePtr;
  Counter counter;
  mlir::Value retValueBox;
};

class AnyVariableStack {
public:
  AnyVariableStack(mlir::Location loc, fir::FirOpBuilder &builder,
                   mlir::Type valueStaticType);
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value);
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
  bool canBeFetchedAfterPush() const { return true; }
 
private:
  mlir::Type variableStaticType;
  mlir::Value opaquePtr;
  Counter counter;
  mlir::Value retValueBox;
};

class AnyAddressStack : public AnyValueStack {
public:
  AnyAddressStack(mlir::Location loc, fir::FirOpBuilder &builder,
                  mlir::Type addressType);
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value);
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
 
private:
  mlir::Type addressType;
};
 
class TemporaryStorage;

class AnyVectorSubscriptStack : public AnyVariableStack {
public:
  AnyVectorSubscriptStack(mlir::Location loc, fir::FirOpBuilder &builder,
                          mlir::Type valueStaticType,
                          bool shapeCanBeSavedAsRegister, int rank);
  void pushShape(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value shape);
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
  mlir::Value fetchShape(mlir::Location loc, fir::FirOpBuilder &builder);
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
  bool canBeFetchedAfterPush() const { return true; }
 
private:
  std::unique_ptr<TemporaryStorage> shapeTemp;
  // If the shape is saved inside a descriptor (as extents),
  // keep track of the descriptor type.
  std::optional<mlir::Type> boxType;
};

class TemporaryStorage {
public:
  template <typename T>
  TemporaryStorage(T &&impl) : impl{std::forward<T>(impl)} {}
 
  void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Value value, mlir::ValueRange indices = {}) {
    // Only ArrayTemp uses the loop indices; other temps don't take them.
    std::visit(Fortran::common::visitors{
                   [&](ArrayTemp &temp) {
                     temp.pushValue(loc, builder, value, indices);
                   },
                   [&](auto &temp) { temp.pushValue(loc, builder, value); }},
               impl);
  }
  void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder) {
    std::visit([&](auto &temp) { temp.resetFetchPosition(loc, builder); },
               impl);
  }
  mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder,
                    mlir::ValueRange indices = {}) {
    return std::visit(
        Fortran::common::visitors{
            [&](ArrayTemp &temp) { return temp.fetch(loc, builder, indices); },
            [&](auto &temp) { return temp.fetch(loc, builder); }},
        impl);
  }
  void destroy(mlir::Location loc, fir::FirOpBuilder &builder) {
    std::visit([&](auto &temp) { temp.destroy(loc, builder); }, impl);
  }
  bool canBeFetchedAfterPush() const {
    return std::visit([&](auto &temp) { return temp.canBeFetchedAfterPush(); },
                      impl);
  }
 
  template <typename T>
  T &cast() {
    return std::get<T>(impl);
  }
 
private:
  std::variant<HomogeneousScalarStack, ArrayTemp, SimpleCopy, SSARegister,
               AnyValueStack, AnyVariableStack, AnyVectorSubscriptStack,
               AnyAddressStack>
      impl;
};
} // namespace fir::factory
#endif // FORTRAN_OPTIMIZER_BUILDER_TEMPORARYSTORAGE_H