big refactor

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

@@ -25,6 +25,9 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
   const int64_t laneCount = physicalType.getDimSize(0);
   if (laneCount <= 0)
     return emitError(loc, "fragment assembly requires at least one physical source slot"), failure();
+  auto physicalLayoutValue = spatial::symbolizePhysicalLayout(physicalLayout);
+  if (!physicalLayoutValue)
+    return emitError(loc, "unknown physical layout for fragment assembly"), failure();
   const int64_t fragmentElements = physicalType.getNumElements() / laneCount;
   SmallVector<int64_t> operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes,
     strides(entries.size() * rank, 1);
@@ -47,13 +50,18 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
     llvm::append_range(offsets, entry.destinationOffsets);
     llvm::append_range(sizes, entry.sizes);
   }
-  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
-    rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
+  auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
+    spatial::getNCHWLayout(rewriter.getContext()),
+    spatial::PhysicalLayoutAttr::get(rewriter.getContext(), *physicalLayoutValue),
     rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
-    rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getStringAttr(indexMap), spatial::getFragmentAssemblyMode(rewriter.getContext()),
     rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
     rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
-    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete"));
+  if (indexMap == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    blueprint.setIndexMapAttr(rewriter.getStringAttr("fragment_assembly"));
+  return blueprint.getOutput();
 }
 
 Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {

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

@@ -394,6 +394,39 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
     rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }
 
+template <typename BodyFn>
+mlir::FailureOr<mlir::Value> mapGraphBatchFragments(mlir::Value input,
+                                                    mlir::RankedTensorType outputType,
+                                                    mlir::PatternRewriter& rewriter,
+                                                    mlir::Location loc,
+                                                    BodyFn&& build) {
+  auto inputType = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
+  if (!inputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != outputType.getRank() || inputType.getRank() < 2
+      || inputType.getDimSize(0) != outputType.getDimSize(0))
+    return mlir::failure();
+  auto inputFragmentType = mlir::RankedTensorType::get(
+    inputType.getShape().drop_front(), inputType.getElementType(), inputType.getEncoding());
+  auto outputFragmentType = mlir::RankedTensorType::get(
+    outputType.getShape().drop_front(), outputType.getElementType(), outputType.getEncoding());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, mlir::TypeRange {outputType}, inputType.getDimSize(0), {}, mlir::ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) -> mlir::LogicalResult {
+      auto fragment = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), args.lane, inputFragmentType);
+      if (mlir::failed(fragment))
+        return mlir::failure();
+      mlir::FailureOr<mlir::Value> result = build(*fragment, outputFragmentType);
+      if (mlir::failed(result) || result->getType() != outputFragmentType)
+        return mlir::failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, *result, args.outputs.front(), args.lane);
+      return mlir::success();
+    });
+  if (mlir::failed(batch))
+    return mlir::failure();
+  return batch->getResult(0);
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -0,0 +1,42 @@
+#include "ContractionPlanning.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+#include <algorithm>
+
+namespace onnx_mlir {
+
+namespace {
+
+static int64_t ceilDivide(int64_t value, int64_t divisor) {
+  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+}
+
+} // namespace
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetResources& target,
+    ContractionPlanKind kind,
+    int64_t laneCount,
+    int64_t fragmentRows) {
+  ContractionPlan plan;
+  plan.tileK = std::max<int64_t>(1, target.matrixShape.rows);
+  plan.tileN = std::max<int64_t>(1, target.matrixShape.columns);
+  plan.reductionSlices = std::max<int64_t>(1, ceilDivide(problem.k, plan.tileK));
+  plan.outputTiles = std::max<int64_t>(1, ceilDivide(problem.n, plan.tileN));
+  const int64_t rowsPerLane = std::max<int64_t>(
+      1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows);
+
+  if (laneCount != 0)
+    plan.laneCount = laneCount;
+  else if (kind == ContractionPlanKind::StaticTiled)
+    plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles;
+  else if (kind == ContractionPlanKind::GroupedRowDynamicVVD)
+    plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane);
+  else
+    plan.laneCount = problem.batch * problem.m * problem.n;
+  return plan;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "ContractionProblem.hpp"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+namespace onnx_mlir {
+
+enum class ContractionPlanKind {
+  StaticTiled,
+  BatchedDynamicVVD,
+  GroupedRowDynamicVVD,
+};
+
+struct ContractionPlan {
+  int64_t tileK = 1;
+  int64_t tileN = 1;
+  int64_t reductionSlices = 1;
+  int64_t outputTiles = 1;
+  int64_t laneCount = 0;
+};
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetResources& target,
+    ContractionPlanKind kind,
+    int64_t laneCount = 0,
+    int64_t fragmentRows = 0);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+struct ContractionProblem {
+  llvm::SmallVector<int64_t> lhsBatchShape;
+  llvm::SmallVector<int64_t> rhsBatchShape;
+  llvm::SmallVector<int64_t> outputBatchShape;
+  int64_t lhsBatch = 1;
+  int64_t rhsBatch = 1;
+  int64_t batch = 1;
+  int64_t m = 0;
+  int64_t k = 0;
+  int64_t n = 0;
+  mlir::Type lhsElementType;
+  mlir::Type rhsElementType;
+  mlir::Type resultElementType;
+};
+
+} // namespace onnx_mlir

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

@@ -1,15 +1,70 @@
 #include "MatrixProductLowering.hpp"
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
+static bool isInsideSpatialCompute(Operation* op) {
+  for (Operation* parent = op; parent; parent = parent->getParentOp())
+    if (spatial::isAnySpatialComputeLike(parent))
+      return true;
+  return false;
+}
+
+static Value buildLinalgTranspose(Value value,
+                                  RankedTensorType resultType,
+                                  ArrayRef<int64_t> permutation,
+                                  PatternRewriter& rewriter,
+                                  Location loc) {
+  Value init = tensor::EmptyOp::create(
+      rewriter, loc, resultType.getShape(), resultType.getElementType());
+  return linalg::TransposeOp::create(
+    rewriter, loc, value, init, permutation).getResult()[0];
+}
+
+static Value materializeConstantTranspose(Value value,
+                                          RankedTensorType resultType,
+                                          ArrayRef<int64_t> permutation,
+                                          PatternRewriter& rewriter) {
+  auto denseAttr = getHostConstDenseElementsAttr(value);
+  if (!denseAttr)
+    return {};
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return {};
+  return getOrCreateConstant(
+      rewriter, rewriter.getInsertionBlock()->getParentOp(), *transposedAttr, resultType);
+}
+
+Value createLinalgTranspose(Value value,
+                            RankedTensorType resultType,
+                            ArrayRef<int64_t> permutation,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  if (Value constant = materializeConstantTranspose(value, resultType, permutation, rewriter))
+    return constant;
+
+  if (isInsideSpatialCompute(rewriter.getInsertionBlock()->getParentOp()))
+    return buildLinalgTranspose(value, resultType, permutation, rewriter, loc);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, loc, TypeRange {resultType}, {}, ValueRange {value},
+    [&](Value input) {
+      spatial::SpatYieldOp::create(
+        rewriter, loc, buildLinalgTranspose(input, resultType, permutation, rewriter, loc));
+    });
+  return compute.getResult(0);
+}
+
 Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));

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

@@ -5,8 +5,16 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
+mlir::Value createLinalgTranspose(mlir::Value value,
+                                  mlir::RankedTensorType resultType,
+                                  llvm::ArrayRef<int64_t> permutation,
+                                  mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc);
+
 mlir::Value createZeroPaddedTensor(mlir::Value value,
                                    mlir::RankedTensorType resultType,
                                    mlir::PatternRewriter& rewriter,

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

@@ -5,9 +5,9 @@
 #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/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
 
 #include <numeric>
 
@@ -33,6 +33,16 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
     tilesPerRow};
 }
 
+FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(Value value) {
+  auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>();
+  auto logicalType = dyn_cast<RankedTensorType>(value.getType());
+  if (!blueprint || !logicalType || blueprint.getOutput() != value
+      || blueprint.getPhysicalLayout() != spatial::PhysicalLayout::NHWCRowStrip
+      || !spatial::isPhysicalView(blueprint.getMode()))
+    return failure();
+  return describeRowStripPhysicalValue(blueprint.getInput(), logicalType);
+}
+
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
   return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
                                 logicalType.getDimSize(1)},
@@ -144,6 +154,35 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> createRowStripStorageBlueprint(Value storage,
+                                                RankedTensorType logicalType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return failure();
+
+  auto blueprint = spatial::SpatBlueprintOp::create(
+      rewriter,
+      loc,
+      logicalType,
+      storage,
+      ValueRange {},
+      spatial::getNCHWLayout(rewriter.getContext()),
+      spatial::getNHWCRowStripLayout(rewriter.getContext()),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getStringAttr(kRowStripIndexMap),
+      spatial::getPhysicalViewMode(rewriter.getContext()),
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr);
+  return blueprint.getOutput();
+}
+
 FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
@@ -160,8 +199,8 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
       if (failed(fragment))
         return failure();
-      Value nchw = ONNXTransposeOp::create(
-        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+      Value nchw = createLinalgTranspose(
+        *fragment, nchwFragmentType, {0, 3, 1, 2}, rewriter, loc);
       publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
       return success();
     });
@@ -176,7 +215,7 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
          value.logicalType.getDimSize(3)}});
     }
-  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
+  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "dense_nchw",
                                          kRowStripIndexMap, rewriter, loc);
 }
 
@@ -186,25 +225,9 @@ static FailureOr<Value> applyRowStripActivation(const RowStripPhysicalValue& val
                                                 Location loc,
                                                 BuildActivation buildActivation) {
   auto storageType = cast<RankedTensorType>(value.storage.getType());
-  const int64_t laneCount = storageType.getDimSize(0);
-  auto batchOp = createSpatComputeBatch(rewriter,
-                                        loc,
-                                        TypeRange {storageType},
-                                        laneCount,
-                                        {},
-                                        ValueRange {value.storage},
-                                        [&](detail::SpatComputeBatchBodyArgs args) {
-                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
-                                            rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
-                                          if (failed(fragment)) return failure();
-                                          Value result = buildActivation(*fragment);
-                                          publishGraphBatchPhysicalFragment(
-                                            rewriter, loc, result, args.outputs.front(), args.lane);
-                                          return success();
-                                        });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return mapGraphBatchFragments(value.storage, storageType, rewriter, loc, [&](Value fragment, RankedTensorType) {
+    return FailureOr<Value>(buildActivation(fragment));
+  });
 }
 
 FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {

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

@@ -6,6 +6,12 @@
 
 namespace onnx_mlir {
 
+namespace spatial {
+class SpatBlueprintOp;
+class SpatFlattenPlanOp;
+struct SpatialTargetResources;
+} // namespace spatial
+
 inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
 
 struct RowStripPhysicalValue {
@@ -18,6 +24,8 @@ struct RowStripPhysicalValue {
 mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
                                                                      mlir::RankedTensorType logicalType);
 
+mlir::FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(mlir::Value value);
+
 std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
 buildRowStripMetadata(mlir::RankedTensorType type);
 
@@ -53,6 +61,11 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> createRowStripStorageBlueprint(mlir::Value storage,
+                                                            mlir::RankedTensorType logicalType,
+                                                            mlir::PatternRewriter& rewriter,
+                                                            mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);
@@ -80,4 +93,14 @@ mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysical
                                                  mlir::PatternRewriter& rewriter,
                                                  mlir::Location loc);
 
+mlir::LogicalResult canLowerFlattenFromRowStrip(
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target);
+
+mlir::LogicalResult lowerFlattenFromRowStrip(
+    const RowStripPhysicalValue& input,
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target,
+    mlir::PatternRewriter& rewriter);
+
 } // namespace onnx_mlir

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

@@ -5,7 +5,6 @@
 
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 
 using namespace mlir;
@@ -67,11 +66,15 @@ sliceVector(const Value& vectorToSlice, int64_t sliceSize, PatternRewriter& rewr
 }
 
 DenseMap<CoreId, SmallVector<Value>>
-sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) {
-  SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice,
+                              PatternRewriter& rewriter,
+                              Location loc,
+                              const spatial::SpatialTargetResources& target) {
+  SmallVector<Value> slices = sliceVector(
+    vectorToSlice, static_cast<int64_t>(target.matrixShape.rows), rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
   for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
-    size_t coreId = sliceId / crossbarCountInCore;
+    size_t coreId = sliceId / target.matrixUnitsPerProcessor;
     slicesPerCore[coreId].push_back(slices[sliceId]);
   }
   return slicesPerCore;

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

@@ -7,6 +7,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -26,6 +27,9 @@ llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
 /// Partitions one logical vector into per-core crossbar-sized slices using the
 /// current PIM target geometry.
 llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
-  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
+  const mlir::Value& vectorToSlice,
+  mlir::PatternRewriter& rewriter,
+  mlir::Location loc,
+  const spatial::SpatialTargetResources& target);
 
 } // namespace onnx_mlir