[Flang] Upstream conversion of the XRebox Op

kiranchandramohan · kiranchandramohan · commit fa517555420b · 2021-12-07T22:22:39.000Z
The XRebox Op is formed by the codegen rewrite which makes it easier to convert the operation to LLVM. The XRebox op includes the information from the rebox op and the associated slice, shift, and shape ops. During the conversion process a new descriptor is created for reboxing. Co-authored-by: Jean Perier <jperier@nvidia.com> Co-authored-by: Eric Schweitz <eschweitz@nvidia.com> Co-authored-by: Val Donaldson <vdonaldson@nvidia.com> Reviewed By: clementval Differential Revision: https://reviews.llvm.org/D114709
diff --git a/flang/include/flang/Optimizer/CodeGen/CGOps.td b/flang/include/flang/Optimizer/CodeGen/CGOps.td
@@ -127,6 +127,14 @@ def fircg_XReboxOp : fircg_Op<"ext_rebox", [AttrSizedOperandSegments]> {
     unsigned getRank();
     // The rank of the result box
     unsigned getOutRank();
+
+    unsigned shapeOffset() { return 1; }
+    unsigned shiftOffset() { return shapeOffset() + shape().size(); }
+    unsigned sliceOffset() { return shiftOffset() + shift().size(); }
+    unsigned subcomponentOffset() { return sliceOffset() + slice().size(); }
+    unsigned substrOffset() {
+      return subcomponentOffset() + subcomponent().size();
+    }
   }];
 }
 
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -38,6 +38,10 @@ static constexpr unsigned defaultAlign = 8;
 static constexpr unsigned kAttrPointer = CFI_attribute_pointer;
 static constexpr unsigned kAttrAllocatable = CFI_attribute_allocatable;
 
+static mlir::Type getVoidPtrType(mlir::MLIRContext *context) {
+  return mlir::LLVM::LLVMPointerType::get(mlir::IntegerType::get(context, 8));
+}
+
 static mlir::LLVM::ConstantOp
 genConstantIndex(mlir::Location loc, mlir::Type ity,
                  mlir::ConversionPatternRewriter &rewriter,
@@ -1854,6 +1858,222 @@ struct XEmboxOpConversion : public EmboxCommonConversion<fir::cg::XEmboxOp> {
   }
 };
 
+/// Create a new box given a box reference.
+struct XReboxOpConversion : public EmboxCommonConversion<fir::cg::XReboxOp> {
+  using EmboxCommonConversion::EmboxCommonConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::cg::XReboxOp rebox, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Location loc = rebox.getLoc();
+    mlir::Type idxTy = lowerTy().indexType();
+    mlir::Value loweredBox = adaptor.getOperands()[0];
+    mlir::ValueRange operands = adaptor.getOperands();
+
+    // Create new descriptor and fill its non-shape related data.
+    llvm::SmallVector<mlir::Value, 2> lenParams;
+    mlir::Type inputEleTy = getInputEleTy(rebox);
+    if (auto charTy = inputEleTy.dyn_cast<fir::CharacterType>()) {
+      mlir::Value len =
+          loadElementSizeFromBox(loc, idxTy, loweredBox, rewriter);
+      if (charTy.getFKind() != 1) {
+        mlir::Value width =
+            genConstantIndex(loc, idxTy, rewriter, charTy.getFKind());
+        len = rewriter.create<mlir::LLVM::SDivOp>(loc, idxTy, len, width);
+      }
+      lenParams.emplace_back(len);
+    } else if (auto recTy = inputEleTy.dyn_cast<fir::RecordType>()) {
+      if (recTy.getNumLenParams() != 0)
+        TODO(loc, "reboxing descriptor of derived type with length parameters");
+    }
+    auto [boxTy, dest, eleSize] =
+        consDescriptorPrefix(rebox, rewriter, rebox.getOutRank(), lenParams);
+
+    // Read input extents, strides, and base address
+    llvm::SmallVector<mlir::Value> inputExtents;
+    llvm::SmallVector<mlir::Value> inputStrides;
+    const unsigned inputRank = rebox.getRank();
+    for (unsigned i = 0; i < inputRank; ++i) {
+      mlir::Value dim = genConstantIndex(loc, idxTy, rewriter, i);
+      SmallVector<mlir::Value, 3> dimInfo =
+          getDimsFromBox(loc, {idxTy, idxTy, idxTy}, loweredBox, dim, rewriter);
+      inputExtents.emplace_back(dimInfo[1]);
+      inputStrides.emplace_back(dimInfo[2]);
+    }
+
+    mlir::Type baseTy = getBaseAddrTypeFromBox(loweredBox.getType());
+    mlir::Value baseAddr =
+        loadBaseAddrFromBox(loc, baseTy, loweredBox, rewriter);
+
+    if (!rebox.slice().empty() || !rebox.subcomponent().empty())
+      return sliceBox(rebox, dest, baseAddr, inputExtents, inputStrides,
+                      operands, rewriter);
+    return reshapeBox(rebox, dest, baseAddr, inputExtents, inputStrides,
+                      operands, rewriter);
+  }
+
+private:
+  /// Write resulting shape and base address in descriptor, and replace rebox
+  /// op.
+  mlir::LogicalResult
+  finalizeRebox(fir::cg::XReboxOp rebox, mlir::Value dest, mlir::Value base,
+                mlir::ValueRange lbounds, mlir::ValueRange extents,
+                mlir::ValueRange strides,
+                mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::Location loc = rebox.getLoc();
+    mlir::Value one = genConstantIndex(loc, lowerTy().indexType(), rewriter, 1);
+    for (auto iter : llvm::enumerate(llvm::zip(extents, strides))) {
+      unsigned dim = iter.index();
+      mlir::Value lb = lbounds.empty() ? one : lbounds[dim];
+      dest = insertLowerBound(rewriter, loc, dest, dim, lb);
+      dest = insertExtent(rewriter, loc, dest, dim, std::get<0>(iter.value()));
+      dest = insertStride(rewriter, loc, dest, dim, std::get<1>(iter.value()));
+    }
+    dest = insertBaseAddress(rewriter, loc, dest, base);
+    mlir::Value result =
+        placeInMemoryIfNotGlobalInit(rewriter, rebox.getLoc(), dest);
+    rewriter.replaceOp(rebox, result);
+    return success();
+  }
+
+  // Apply slice given the base address, extents and strides of the input box.
+  mlir::LogicalResult
+  sliceBox(fir::cg::XReboxOp rebox, mlir::Value dest, mlir::Value base,
+           mlir::ValueRange inputExtents, mlir::ValueRange inputStrides,
+           mlir::ValueRange operands,
+           mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::Location loc = rebox.getLoc();
+    mlir::Type voidPtrTy = ::getVoidPtrType(rebox.getContext());
+    mlir::Type idxTy = lowerTy().indexType();
+    mlir::Value zero = genConstantIndex(loc, idxTy, rewriter, 0);
+    // Apply subcomponent and substring shift on base address.
+    if (!rebox.subcomponent().empty() || !rebox.substr().empty()) {
+      // Cast to inputEleTy* so that a GEP can be used.
+      mlir::Type inputEleTy = getInputEleTy(rebox);
+      auto llvmElePtrTy =
+          mlir::LLVM::LLVMPointerType::get(convertType(inputEleTy));
+      base = rewriter.create<mlir::LLVM::BitcastOp>(loc, llvmElePtrTy, base);
+
+      if (!rebox.subcomponent().empty()) {
+        llvm::SmallVector<mlir::Value> gepOperands = {zero};
+        for (unsigned i = 0; i < rebox.subcomponent().size(); ++i)
+          gepOperands.push_back(operands[rebox.subcomponentOffset() + i]);
+        base = genGEP(loc, llvmElePtrTy, rewriter, base, gepOperands);
+      }
+      if (!rebox.substr().empty())
+        base = shiftSubstringBase(rewriter, loc, base,
+                                  operands[rebox.substrOffset()]);
+    }
+
+    if (rebox.slice().empty())
+      // The array section is of the form array[%component][substring], keep
+      // the input array extents and strides.
+      return finalizeRebox(rebox, dest, base, /*lbounds*/ llvm::None,
+                           inputExtents, inputStrides, rewriter);
+
+    // Strides from the fir.box are in bytes.
+    base = rewriter.create<mlir::LLVM::BitcastOp>(loc, voidPtrTy, base);
+
+    // The slice is of the form array(i:j:k)[%component]. Compute new extents
+    // and strides.
+    llvm::SmallVector<mlir::Value> slicedExtents;
+    llvm::SmallVector<mlir::Value> slicedStrides;
+    mlir::Value one = genConstantIndex(loc, idxTy, rewriter, 1);
+    const bool sliceHasOrigins = !rebox.shift().empty();
+    unsigned sliceOps = rebox.sliceOffset();
+    unsigned shiftOps = rebox.shiftOffset();
+    auto strideOps = inputStrides.begin();
+    const unsigned inputRank = inputStrides.size();
+    for (unsigned i = 0; i < inputRank;
+         ++i, ++strideOps, ++shiftOps, sliceOps += 3) {
+      mlir::Value sliceLb =
+          integerCast(loc, rewriter, idxTy, operands[sliceOps]);
+      mlir::Value inputStride = *strideOps; // already idxTy
+      // Apply origin shift: base += (lb-shift)*input_stride
+      mlir::Value sliceOrigin =
+          sliceHasOrigins
+              ? integerCast(loc, rewriter, idxTy, operands[shiftOps])
+              : one;
+      mlir::Value diff =
+          rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, sliceLb, sliceOrigin);
+      mlir::Value offset =
+          rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, diff, inputStride);
+      base = genGEP(loc, voidPtrTy, rewriter, base, offset);
+      // Apply upper bound and step if this is a triplet. Otherwise, the
+      // dimension is dropped and no extents/strides are computed.
+      mlir::Value upper = operands[sliceOps + 1];
+      const bool isTripletSlice =
+          !mlir::isa_and_nonnull<mlir::LLVM::UndefOp>(upper.getDefiningOp());
+      if (isTripletSlice) {
+        mlir::Value step =
+            integerCast(loc, rewriter, idxTy, operands[sliceOps + 2]);
+        // extent = ub-lb+step/step
+        mlir::Value sliceUb = integerCast(loc, rewriter, idxTy, upper);
+        mlir::Value extent = computeTripletExtent(rewriter, loc, sliceLb,
+                                                  sliceUb, step, zero, idxTy);
+        slicedExtents.emplace_back(extent);
+        // stride = step*input_stride
+        mlir::Value stride =
+            rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, step, inputStride);
+        slicedStrides.emplace_back(stride);
+      }
+    }
+    return finalizeRebox(rebox, dest, base, /*lbounds*/ llvm::None,
+                         slicedExtents, slicedStrides, rewriter);
+  }
+
+  /// Apply a new shape to the data described by a box given the base address,
+  /// extents and strides of the box.
+  mlir::LogicalResult
+  reshapeBox(fir::cg::XReboxOp rebox, mlir::Value dest, mlir::Value base,
+             mlir::ValueRange inputExtents, mlir::ValueRange inputStrides,
+             mlir::ValueRange operands,
+             mlir::ConversionPatternRewriter &rewriter) const {
+    mlir::ValueRange reboxShifts{operands.begin() + rebox.shiftOffset(),
+                                 operands.begin() + rebox.shiftOffset() +
+                                     rebox.shift().size()};
+    if (rebox.shape().empty()) {
+      // Only setting new lower bounds.
+      return finalizeRebox(rebox, dest, base, reboxShifts, inputExtents,
+                           inputStrides, rewriter);
+    }
+
+    mlir::Location loc = rebox.getLoc();
+    // Strides from the fir.box are in bytes.
+    mlir::Type voidPtrTy = ::getVoidPtrType(rebox.getContext());
+    base = rewriter.create<mlir::LLVM::BitcastOp>(loc, voidPtrTy, base);
+
+    llvm::SmallVector<mlir::Value> newStrides;
+    llvm::SmallVector<mlir::Value> newExtents;
+    mlir::Type idxTy = lowerTy().indexType();
+    // First stride from input box is kept. The rest is assumed contiguous
+    // (it is not possible to reshape otherwise). If the input is scalar,
+    // which may be OK if all new extents are ones, the stride does not
+    // matter, use one.
+    mlir::Value stride = inputStrides.empty()
+                             ? genConstantIndex(loc, idxTy, rewriter, 1)
+                             : inputStrides[0];
+    for (unsigned i = 0; i < rebox.shape().size(); ++i) {
+      mlir::Value rawExtent = operands[rebox.shapeOffset() + i];
+      mlir::Value extent = integerCast(loc, rewriter, idxTy, rawExtent);
+      newExtents.emplace_back(extent);
+      newStrides.emplace_back(stride);
+      // nextStride = extent * stride;
+      stride = rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, extent, stride);
+    }
+    return finalizeRebox(rebox, dest, base, reboxShifts, newExtents, newStrides,
+                         rewriter);
+  }
+
+  /// Return scalar element type of the input box.
+  static mlir::Type getInputEleTy(fir::cg::XReboxOp rebox) {
+    auto ty = fir::dyn_cast_ptrOrBoxEleTy(rebox.box().getType());
+    if (auto seqTy = ty.dyn_cast<fir::SequenceType>())
+      return seqTy.getEleTy();
+    return ty;
+  }
+};
+
 // Code shared between insert_value and extract_value Ops.
 struct ValueOpCommon {
   // Translate the arguments pertaining to any multidimensional array to
@@ -2616,7 +2836,8 @@ class FIRToLLVMLowering : public fir::FIRToLLVMLoweringBase<FIRToLLVMLowering> {
         SliceOpConversion, StoreOpConversion, StringLitOpConversion,
         SubcOpConversion, UnboxCharOpConversion, UnboxProcOpConversion,
         UndefOpConversion, UnreachableOpConversion, XArrayCoorOpConversion,
-        XEmboxOpConversion, ZeroOpConversion>(typeConverter);
+        XEmboxOpConversion, XReboxOpConversion, ZeroOpConversion>(
+        typeConverter);
     mlir::populateStdToLLVMConversionPatterns(typeConverter, pattern);
     mlir::arith::populateArithmeticToLLVMConversionPatterns(typeConverter,
                                                             pattern);
diff --git a/flang/test/Fir/convert-to-llvm.fir b/flang/test/Fir/convert-to-llvm.fir
