even faster on pimcomp models

src/PIM/Compiler/CMakeLists.txt

@@ -31,6 +31,7 @@ add_pim_library(OMPimCompilerUtils
   OMPimCommon
   OMPimBufferization
   OMPimHostConstantFolding
+  OMPimInstructionSelection
   OMPimLocalMemoryPlanning
   OMPimVerification
   OMPimPasses

src/PIM/Compiler/PimCodeGen.cpp

@@ -108,6 +108,23 @@ static Operation* getDiagnosticAnchor(mlir::Value value) {
   return nullptr;
 }
 
+static bool isZeroSplatGlobal(mlir::Value value) {
+  auto getGlobalOp = value.getDefiningOp<memref::GetGlobalOp>();
+  auto moduleOp = getGlobalOp ? getGlobalOp->getParentOfType<ModuleOp>() : ModuleOp();
+  auto globalOp = moduleOp ? lookupGlobalForGetGlobal(moduleOp, getGlobalOp) : memref::GlobalOp();
+  if (!globalOp || !globalOp.getConstant() || !globalOp.getInitialValue())
+    return false;
+  auto denseAttr = dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue());
+  if (!denseAttr || !denseAttr.isSplat())
+    return false;
+  Attribute valueAttr = denseAttr.getSplatValue<Attribute>();
+  if (auto floatAttr = dyn_cast<FloatAttr>(valueAttr))
+    return floatAttr.getValue().isZero() && !floatAttr.getValue().isNegative();
+  if (auto integerAttr = dyn_cast<IntegerAttr>(valueAttr))
+    return integerAttr.getValue().isZero();
+  return false;
+}
+
 // PIM instruction immediates are serialized as signed int32_t fields today
 // (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
 // the non-negative int32_t range.
@@ -541,14 +558,23 @@ void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
                                size_t rs1Offset,
                                size_t size,
                                StringRef sizeFieldName) const {
-  setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);
-
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic1 = 0;
-  instruction.generic2 = 0;
+  if (rdOffset == rs1Offset) {
+    setupRdRs1(rdAddr, 0, rs1Addr, 0);
+    instruction.generic1 = rdOffset == 0 ? 0 : 3;
+    instruction.generic2 = pim::checkedI32OrCrash(rdOffset, "shared address offset");
+  } else if (rdOffset != 0) {
+    setupRdRs1(rdAddr, 0, rs1Addr, rs1Offset);
+    instruction.generic1 = 1;
+    instruction.generic2 = pim::checkedI32OrCrash(rdOffset, "rd address offset");
+  } else {
+    setupRdRs1(rdAddr, 0, rs1Addr, 0);
+    instruction.generic1 = rs1Offset == 0 ? 0 : 2;
+    instruction.generic2 = pim::checkedI32OrCrash(rs1Offset, "rs1 address offset");
+  }
   instruction.generic3 = pim::checkedI32OrCrash(size, sizeFieldName);
   emitInstruction(instruction);
 }
@@ -584,6 +610,16 @@ void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticVa
   auto hostSourceOffset = indexOf(loadOp.getHostSourceOffset(), knowledge);
   assert(succeeded(deviceTargetOffset) && succeeded(hostSourceOffset)
          && "pim.memcp_hd offsets must be statically resolvable during codegen");
+  if (isZeroSplatGlobal(loadOp.getHostSource())) {
+    setupRd(addressOf(loadOp.getDeviceTarget(), knowledge), *deviceTargetOffset);
+    pim_binary::InstructionRecord instruction;
+    instruction.opcode = pim_binary::Opcode::lldi;
+    instruction.rd = 0;
+    instruction.r2OrImm = 0;
+    instruction.generic3 = loadOp.getSize();
+    emitInstruction(instruction);
+    return;
+  }
   emitMemCopyOp(pim_binary::Opcode::ld,
                 addressOf(loadOp.getDeviceTarget(), knowledge),
                 *deviceTargetOffset,
@@ -619,6 +655,26 @@ void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledg
                 "len");
 }
 
+void PimCodeGen::codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& knowledge) const {
+  auto targetOffset = indexOf(vmvOp.getTargetOffset(), knowledge);
+  auto sourceOffset = indexOf(vmvOp.getSourceOffset(), knowledge);
+  auto sourceStride = indexOf(vmvOp.getSourceStride(), knowledge);
+  assert(succeeded(targetOffset) && succeeded(sourceOffset) && succeeded(sourceStride)
+         && "pim.vmv operands must be statically resolvable during codegen");
+  auto sourceType = cast<ShapedType>(vmvOp.getSource().getType());
+  int32_t bitwidth = getVectorElementBitwidthOrCrash(sourceType);
+  ensureVectorBitwidth(bitwidth, bitwidth);
+  setupRdRs1Rs2(addressOf(vmvOp.getTarget(), knowledge), *targetOffset,
+                addressOf(vmvOp.getSource(), knowledge), *sourceOffset, 0, *sourceStride);
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vmv;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = vmvOp.getLength();
+  emitInstruction(instruction);
+}
+
 void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const {
   auto outputOffset = indexOf(receiveOp.getOutputOffset(), knowledge);
   auto sourceCoreId = indexOf(receiveOp.getSourceCoreId(), knowledge);
@@ -723,89 +779,6 @@ void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode,
   emitInstruction(instruction);
 }
 
-void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const {
-  auto srcAddr = addressOf(plan.source, knowledge);
-  auto dstAddr = addressOf(plan.destination, knowledge);
-
-  size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes;
-  int32_t maxSourceAddress = pim::checkedI32OrCrash(
-    pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address");
-  int32_t maxDestinationAddress =
-    pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
-                           "transpose destination address");
-  (void) maxSourceAddress;
-  (void) maxDestinationAddress;
-
-  pim_binary::InstructionRecord copyInstruction;
-  copyInstruction.opcode = pim_binary::Opcode::lmv;
-  copyInstruction.rd = 0;
-  copyInstruction.r1 = 1;
-
-  size_t maxRunElements = static_cast<size_t>(std::numeric_limits<int32_t>::max()) / plan.elementBytes;
-  auto emitRun = [&](size_t sourceStart, size_t destinationStart, size_t runLength) {
-    while (runLength != 0) {
-      size_t chunkElements = std::min(runLength, maxRunElements);
-      // totalBytes was checked when the plan was compiled, so these bounded
-      // products cannot overflow.
-      size_t sourceOffset = sourceStart * plan.elementBytes;
-      size_t destinationOffset = destinationStart * plan.elementBytes;
-      size_t byteSize = chunkElements * plan.elementBytes;
-      assert(sourceOffset <= plan.totalBytes - byteSize && destinationOffset <= plan.totalBytes - byteSize);
-      genSetRegisterImmediate(0, static_cast<int32_t>(dstAddr + destinationOffset));
-      genSetRegisterImmediate(1, static_cast<int32_t>(srcAddr + sourceOffset));
-      copyInstruction.generic3 = static_cast<int32_t>(byteSize);
-      emitInstruction(copyInstruction);
-      sourceStart += chunkElements;
-      destinationStart += chunkElements;
-      runLength -= chunkElements;
-    }
-  };
-
-  if (plan.storagePreserving) {
-    emitRun(0, 0, plan.totalElements);
-    return;
-  }
-
-  size_t rank = plan.sourceShape.size();
-  SmallVector<size_t> sourceIndices(rank, 0);
-  size_t destinationFlat = 0;
-  size_t runSourceStart = 0;
-  size_t runDestinationStart = 0;
-  size_t runLength = 0;
-
-  for (size_t sourceFlat = 0; sourceFlat < plan.totalElements; ++sourceFlat) {
-    if (runLength != 0 && destinationFlat != runDestinationStart + runLength) {
-      emitRun(runSourceStart, runDestinationStart, runLength);
-      runSourceStart = sourceFlat;
-      runDestinationStart = destinationFlat;
-      runLength = 0;
-    }
-    if (runLength == 0) {
-      runSourceStart = sourceFlat;
-      runDestinationStart = destinationFlat;
-    }
-    ++runLength;
-
-    if (runLength == maxRunElements) {
-      emitRun(runSourceStart, runDestinationStart, runLength);
-      runLength = 0;
-    }
-
-    if (sourceFlat + 1 == plan.totalElements)
-      break;
-    for (size_t sourceDim = rank; sourceDim-- > 0;) {
-      destinationFlat += plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]];
-      if (++sourceIndices[sourceDim] < plan.sourceShape[sourceDim])
-        break;
-      sourceIndices[sourceDim] = 0;
-      destinationFlat -= plan.destinationRewinds[sourceDim];
-    }
-  }
-
-  if (runLength != 0)
-    emitRun(runSourceStart, runDestinationStart, runLength);
-}
-
 static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
   SmallVector<Operation*> coreLikeOps;
   for (Operation& op : funcOp.getBody().front())
@@ -1015,6 +988,7 @@ static LogicalResult executeCompiledCorePlan(
       coreCodeGen.codeGenStoreOp(cast<pim::PimMemCopyDevToHostOp>(node.op), knowledge);
       break;
     case CompiledCoreOpKind::Lmv:     coreCodeGen.codeGenLmvOp(cast<pim::PimMemCopyOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::VMV:     coreCodeGen.codeGenVMVOp(cast<pim::PimVMVOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Send:    coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Concat:  coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
@@ -1024,7 +998,6 @@ static LogicalResult executeCompiledCorePlan(
       else
         return failure();
       break;
-    case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
     case CompiledCoreOpKind::VVAdd:     emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
     case CompiledCoreOpKind::VVSub:     emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
     case CompiledCoreOpKind::VVMul:     emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
@@ -1224,7 +1197,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
 
     for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
       xbarsPerGroup.push_back(weights[slot].shape[1] / static_cast<int64_t>(crossbarSize));
-      std::string sourcePath = outputDirPath + "/weights/" + fileName;
+      std::string sourcePath = "../weights/" + fileName;
       std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
       sys::fs::remove(targetPath);
       if (auto error = sys::fs::create_link(sourcePath, targetPath)) {

src/PIM/Compiler/PimCodeGen.hpp

@@ -25,7 +25,6 @@
 namespace onnx_mlir {
 
 struct CompiledCoreProgram;
-struct CompiledTransposePlan;
 class PimInstructionWriter;
 
 struct MemEntry {
@@ -214,6 +213,7 @@ public:
   void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
   void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
   void codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& knowledge) const;
 
   void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
   void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
@@ -222,7 +222,6 @@ public:
   template <typename MVMTy>
   void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);
 
-  void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const;
 };
 
 OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -281,6 +281,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   verifyExplicitPimCoreCount();
 
   if (pimOnlyCodegen) {
+    pm.addPass(createPimInstructionSelectionPass());
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createPimVerificationPass());
     pm.addPass(createEmitPimCodePass());
@@ -315,6 +316,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(mlir::createLowerAffinePass());
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
+    pm.addPass(createPimInstructionSelectionPass());
+    pm.addPass(createMessagePass("Pim instructions selected"));
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createMessagePass("Pim local memory planned"));
     pm.addPass(createPimVerificationPass());

src/PIM/Compiler/PimCoreProgram.cpp

@@ -1,12 +1,6 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/BuiltinTypes.h"
-
-#include "llvm/ADT/STLExtras.h"
-
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -20,11 +14,11 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   if (isa<pim::PimMemCopyHostToDevOp>(op)) return CompiledCoreOpKind::Load;
   if (isa<pim::PimMemCopyDevToHostOp>(op)) return CompiledCoreOpKind::Store;
   if (isa<pim::PimMemCopyOp>(op)) return CompiledCoreOpKind::Lmv;
+  if (isa<pim::PimVMVOp>(op)) return CompiledCoreOpKind::VMV;
   if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
   if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
   if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
   if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
-  if (isa<pim::PimTransposeOp>(op)) return CompiledCoreOpKind::Transpose;
   if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
   if (isa<pim::PimVVSubOp>(op)) return CompiledCoreOpKind::VVSub;
   if (isa<pim::PimVVMulOp>(op)) return CompiledCoreOpKind::VVMul;
@@ -38,77 +32,6 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   return failure();
 }
 
-static bool isStoragePreservingTranspose(ArrayRef<size_t> sourceShape, ArrayRef<int64_t> permutation) {
-  SmallVector<unsigned> sourceNonUnitDims;
-  SmallVector<unsigned> destinationSourceNonUnitDims;
-  for (auto [dim, size] : llvm::enumerate(sourceShape))
-    if (size != 1)
-      sourceNonUnitDims.push_back(dim);
-  for (int64_t sourceDim : permutation)
-    if (sourceShape[sourceDim] != 1)
-      destinationSourceNonUnitDims.push_back(static_cast<unsigned>(sourceDim));
-  return sourceNonUnitDims == destinationSourceNonUnitDims;
-}
-
-static FailureOr<CompiledTransposePlan> compileTransposePlan(pim::PimTransposeOp transposeOp) {
-  auto sourceType = cast<ShapedType>(transposeOp.getInput().getType());
-  ArrayRef<int64_t> sourceShape = sourceType.getShape();
-  size_t rank = sourceShape.size();
-  CompiledTransposePlan plan;
-  plan.source = transposeOp.getInput();
-  plan.destination = transposeOp.getOutputBuffer();
-  plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
-  auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements");
-  if (failed(totalElements)) return failure();
-  plan.totalElements = *totalElements;
-  auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size");
-  if (failed(totalBytes)) return failure();
-  plan.totalBytes = *totalBytes;
-
-  SmallVector<int64_t> permutation = map_to_vector(transposeOp.getPermutation().getAsRange<IntegerAttr>(),
-                                                   [](IntegerAttr attr) { return attr.getInt(); });
-  if (permutation.size() != rank) {
-    transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen");
-    return failure();
-  }
-
-  SmallVector<size_t> destinationShape(rank);
-  plan.destinationStrides.assign(rank, 1);
-  plan.destinationDimensionForSource.assign(rank, 0);
-  plan.destinationRewinds.assign(rank, 0);
-  SmallVector<bool> seenSourceDimensions(rank, false);
-  for (size_t dim = 0; dim < rank; ++dim) {
-    auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension");
-    if (failed(size)) return failure();
-    plan.sourceShape.push_back(*size);
-  }
-  for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) {
-    if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seenSourceDimensions[sourceDim]) {
-      transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once");
-      return failure();
-    }
-    seenSourceDimensions[sourceDim] = true;
-    destinationShape[destinationDim] = plan.sourceShape[sourceDim];
-    plan.destinationDimensionForSource[sourceDim] = destinationDim;
-  }
-  for (size_t dim = rank; dim > 1; --dim) {
-    auto stride = pim::checkedMul(
-      plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride");
-    if (failed(stride)) return failure();
-    plan.destinationStrides[dim - 2] = *stride;
-  }
-  for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) {
-    auto rewind = pim::checkedMul(plan.sourceShape[sourceDim],
-                                  plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]],
-                                  transposeOp,
-                                  "transpose destination rewind");
-    if (failed(rewind)) return failure();
-    plan.destinationRewinds[sourceDim] = *rewind;
-  }
-  plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation);
-  return plan;
-}
-
 static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<CompiledCoreNode>& plan) {
   for (Operation& op : block) {
     if (isa<pim::PimHaltOp, scf::YieldOp, memref::GetGlobalOp>(op) || isCoreStaticAddressOp(&op))
@@ -188,11 +111,6 @@ static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<Compi
     CompiledCoreNode node;
     node.op = &op;
     node.opKind = *opKind;
-    if (*opKind == CompiledCoreOpKind::Transpose) {
-      auto transposePlan = compileTransposePlan(cast<pim::PimTransposeOp>(op));
-      if (failed(transposePlan)) return failure();
-      node.transposePlan = *transposePlan;
-    }
     plan.push_back(std::move(node));
   }
   return success();

src/PIM/Compiler/PimCoreProgram.hpp

@@ -6,34 +6,19 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <memory>
-#include <optional>
-
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 
 namespace onnx_mlir {
 
-struct CompiledTransposePlan {
-  mlir::Value source;
-  mlir::Value destination;
-  size_t elementBytes = 0;
-  size_t totalElements = 0;
-  size_t totalBytes = 0;
-  llvm::SmallVector<size_t> sourceShape;
-  llvm::SmallVector<size_t> destinationStrides;
-  llvm::SmallVector<unsigned> destinationDimensionForSource;
-  llvm::SmallVector<size_t> destinationRewinds;
-  bool storagePreserving = false;
-};
-
 enum class CompiledCoreOpKind : uint8_t {
   Load,
   Store,
   Lmv,
+  VMV,
   Receive,
   Send,
   Concat,
   Vmm,
-  Transpose,
   VVAdd,
   VVSub,
   VVMul,
@@ -62,7 +47,6 @@ struct CompiledCoreNode {
   llvm::SmallVector<int64_t> caseValues;
   llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
-  std::optional<CompiledTransposePlan> transposePlan;
 };
 
 struct CompiledCoreProgram {