meno diamantini

src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.cpp

@@ -0,0 +1,112 @@
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
+  if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  if (resultType.getRank() != 4)
+    return failure();
+  if (biasType.getElementType() != resultType.getElementType())
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  ArrayRef<int64_t> shape = biasType.getShape();
+  if (shape.empty())
+    return success();
+  if (shape.size() == 1)
+    return success(shape[0] == channels);
+  if (shape.size() == 2)
+    return success(shape[0] == 1 && shape[1] == channels);
+  if (shape.size() == 4)
+    return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
+  return failure();
+}
+
+FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
+  auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  if (denseAttr.isSplat()) {
+    return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
+  }
+
+  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
+  if (biasType.getRank() == 1)
+    return flattened;
+  if (biasType.getRank() == 2)
+    return flattened;
+
+  SmallVector<Attribute> channelValues;
+  channelValues.reserve(channels);
+  const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
+  for (int64_t channel = 0; channel < channels; ++channel)
+    channelValues.push_back(flattened[channel * channelStride]);
+  return channelValues;
+}
+
+bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
+  auto attr = getHostConstDenseElementsAttr(bias);
+  if (!attr)
+    return false;
+  auto biasType = dyn_cast<RankedTensorType>(attr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return false;
+  if (failed(getBiasChannelValues(attr, resultType)))
+    return false;
+  if (denseAttr)
+    *denseAttr = attr;
+  return true;
+}
+
+FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
+  auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
+  auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
+  if (!lhsType || !rhsType)
+    return failure();
+  if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
+    return BiasAddPlanCandidate {lhs, rhs};
+  if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
+    return BiasAddPlanCandidate {rhs, lhs};
+  return failure();
+}
+
+FailureOr<Value>
+materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
+    return failure();
+
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> resultValues;
+  resultValues.reserve(resultType.getNumElements());
+  const int64_t batches = resultType.getDimSize(0);
+  const int64_t channels = resultType.getDimSize(1);
+  const int64_t height = resultType.getDimSize(2);
+  const int64_t width = resultType.getDimSize(3);
+  for (int64_t n = 0; n < batches; ++n)
+    for (int64_t c = 0; c < channels; ++c)
+      for (int64_t h = 0; h < height; ++h)
+        for (int64_t w = 0; w < width; ++w)
+          resultValues.push_back((*channelValues)[c]);
+
+  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+
+struct BiasAddPlanCandidate {
+  mlir::Value data;
+  mlir::Value bias;
+};
+
+mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType);
+bool isSupportedBiasAddValue(mlir::Value bias,
+                             mlir::RankedTensorType resultType,
+                             mlir::DenseElementsAttr* denseAttr = nullptr);
+mlir::FailureOr<llvm::SmallVector<mlir::Attribute>>
+getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType);
+mlir::FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(mlir::Value lhs,
+                                                                   mlir::Value rhs,
+                                                                   mlir::RankedTensorType resultType);
+mlir::FailureOr<mlir::Value> materializeDenseBiasAddTensor(mlir::Value bias,
+                                                           mlir::RankedTensorType resultType,
+                                                           mlir::RewriterBase& rewriter,
+                                                           mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp

@@ -0,0 +1,239 @@
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
+  SmallVector<int64_t> offsets;
+  SmallVector<int64_t> sizes;
+  const int64_t channels = type.getDimSize(1);
+  const int64_t height = type.getDimSize(2);
+  const int64_t width = type.getDimSize(3);
+  offsets.reserve(height * 4);
+  sizes.reserve(height * 4);
+  for (int64_t row = 0; row < height; ++row) {
+    offsets.append({0, 0, row, 0});
+    sizes.append({1, channels, 1, width});
+  }
+  return {offsets, sizes};
+}
+
+SmallVector<OpFoldResult> buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) {
+  return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+}
+
+SmallVector<OpFoldResult> buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) {
+  return {rewriter.getIndexAttr(1),
+          rewriter.getIndexAttr(logicalType.getDimSize(1)),
+          rewriter.getIndexAttr(1),
+          rewriter.getIndexAttr(logicalType.getDimSize(3))};
+}
+
+Value extractRowStripFragment(Value storage,
+                              RankedTensorType logicalType,
+                              OpFoldResult row,
+                              PatternRewriter& rewriter,
+                              Location loc) {
+  return tensor::ExtractSliceOp::create(rewriter,
+                                        loc,
+                                        getRowStripFragmentType(logicalType),
+                                        storage,
+                                        buildRowStripFragmentOffsets(rewriter, row),
+                                        buildRowStripFragmentSizes(rewriter, logicalType),
+                                        getUnitStrides(rewriter, 4));
+}
+
+void insertRowStripFragment(Value fragment,
+                            Value output,
+                            RankedTensorType logicalType,
+                            OpFoldResult row,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  createParallelInsertSliceIntoBatchOutput(rewriter,
+                                           loc,
+                                           fragment,
+                                           output,
+                                           buildRowStripFragmentOffsets(rewriter, row),
+                                           buildRowStripFragmentSizes(rewriter, logicalType),
+                                           getUnitStrides(rewriter, 4));
+}
+
+FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
+                                                  RankedTensorType fragmentType,
+                                                  PatternRewriter& rewriter) {
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> values;
+  values.reserve(fragmentType.getNumElements());
+  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
+    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
+      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+          values.push_back((*channelValues)[channel]);
+
+  auto attr = DenseElementsAttr::get(fragmentType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
+}
+
+FailureOr<Value> createRowStripStorageFromRows(Value rows,
+                                               RankedTensorType logicalType,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
+  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
+    return failure();
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
+    return failure();
+  if (logicalType.getDimSize(0) != 1)
+    return failure();
+  if (rowsType.getElementType() != logicalType.getElementType())
+    return failure();
+
+  const int64_t channels = logicalType.getDimSize(1);
+  const int64_t height = logicalType.getDimSize(2);
+  const int64_t width = logicalType.getDimSize(3);
+  if (rowsType.getDimSize(0) != height * width)
+    return failure();
+  if (rowsType.getDimSize(1) != channels)
+    return failure();
+
+  auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
+  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
+      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
+      Value rowSlice = tensor::ExtractSliceOp::create(
+        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+      Value channelWidth = ONNXTransposeOp::create(
+        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
+      Value fragment = tensor::ExpandShapeOp::create(
+        rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
+      insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return failure();
+
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+      createParallelInsertSliceIntoBatchOutput(rewriter,
+                                               loc,
+                                               fragment,
+                                               args.outputs.front(),
+                                               SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                                                          rewriter.getIndexAttr(0),
+                                                                          args.lane,
+                                                                          rewriter.getIndexAttr(0)},
+                                               buildRowStripFragmentSizes(rewriter, logicalType),
+                                               getUnitStrides(rewriter, 4));
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
+    return failure();
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          Value constant;
+                                          if (denseAttr.isSplat()) {
+                                            constant = getOrCreateConstant(
+                                              rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
+                                              fragmentType);
+                                          }
+                                          else {
+                                            FailureOr<Value> perChannel =
+                                              createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
+                                            if (failed(perChannel))
+                                              return failure();
+                                            constant = *perChannel;
+                                          }
+                                          fragment =
+                                            spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
+
+struct RowStripPhysicalValue {
+  mlir::Value storage;
+  mlir::RankedTensorType logicalType;
+  llvm::SmallVector<int64_t, 16> fragmentOffsets;
+  llvm::SmallVector<int64_t, 16> fragmentSizes;
+};
+
+std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
+buildRowStripMetadata(mlir::RankedTensorType type);
+
+mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);
+
+mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
+                                                                   mlir::OpFoldResult row);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
+                                                                 mlir::RankedTensorType logicalType);
+
+mlir::Value extractRowStripFragment(mlir::Value storage,
+                                    mlir::RankedTensorType logicalType,
+                                    mlir::OpFoldResult row,
+                                    mlir::PatternRewriter& rewriter,
+                                    mlir::Location loc);
+
+void insertRowStripFragment(mlir::Value fragment,
+                            mlir::Value output,
+                            mlir::RankedTensorType logicalType,
+                            mlir::OpFoldResult row,
+                            mlir::PatternRewriter& rewriter,
+                            mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
+                                                              mlir::RankedTensorType fragmentType,
+                                                              mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
+                                                           mlir::RankedTensorType logicalType,
+                                                           mlir::PatternRewriter& rewriter,
+                                                           mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> materializeRowStripStorageToDense(mlir::Value storage,
+                                                               mlir::RankedTensorType logicalType,
+                                                               mlir::PatternRewriter& rewriter,
+                                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
+                                               mlir::RankedTensorType logicalType,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
+                                                  mlir::RankedTensorType logicalType,
+                                                  mlir::Value bias,
+                                                  mlir::PatternRewriter& rewriter,
+                                                  mlir::Location loc);
+
+} // namespace onnx_mlir