fix bufferization and weight emission after new gemm patterns

src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.cpp

@@ -30,6 +30,15 @@ Value materializeContiguousMemRef(Value memrefValue, Location loc, RewriterBase&
     .getOutput();
 }
 
+Value allocateContiguousMemRefLike(Value memrefValue, Location loc, RewriterBase& rewriter) {
+  if (succeeded(resolveContiguousAddress(memrefValue)))
+    return memrefValue;
+
+  auto shapedType = cast<ShapedType>(memrefValue.getType());
+  auto contiguousType = MemRefType::get(shapedType.getShape(), shapedType.getElementType());
+  return memref::AllocOp::create(rewriter, loc, contiguousType);
+}
+
 FailureOr<Value>
 getBufferOrValue(RewriterBase& rewriter, Value value, const BufferizationOptions& options, BufferizationState& state) {
   if (isa<BufferLikeType>(value.getType()))

src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp

@@ -6,6 +6,7 @@
 namespace onnx_mlir::pim {
 
 mlir::Value materializeContiguousMemRef(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
+mlir::Value allocateContiguousMemRefLike(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
 
 llvm::FailureOr<mlir::Value> getBufferOrValue(mlir::RewriterBase& rewriter,
                                               mlir::Value value,

src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp

@@ -431,8 +431,11 @@ struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeO
     if (failed(outputBufferOpt))
       return failure();
 
+    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
+    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
+
     replaceOpWithNewBufferizedOp<PimTransposeOp>(
-      rewriter, op, outputBufferOpt->getType(), *inputOpt, transposeOp.getPermutation(), *outputBufferOpt);
+      rewriter, op, contiguousOutput.getType(), contiguousInput, transposeOp.getPermutation(), contiguousOutput);
     return success();
   }
 };
@@ -473,9 +476,10 @@ struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface,
       return failure();
 
     Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
+    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
     replaceOpWithNewBufferizedOp<PimVMMOp>(
-      rewriter, op, outputBufferOpt->getType(), *weightOpt, contiguousInput, *outputBufferOpt);
+      rewriter, op, contiguousOutput.getType(), *weightOpt, contiguousInput, contiguousOutput);
     return success();
   }
 };
@@ -512,9 +516,10 @@ struct BinaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<BinaryDstO
 
     Value contiguousLhs = materializeContiguousMemRef(*lhsOpt, op->getLoc(), rewriter);
     Value contiguousRhs = materializeContiguousMemRef(*rhsOpt, op->getLoc(), rewriter);
+    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
     replaceOpWithNewBufferizedOp<OpTy>(
-      rewriter, op, outputBufferOpt->getType(), contiguousLhs, contiguousRhs, *outputBufferOpt);
+      rewriter, op, contiguousOutput.getType(), contiguousLhs, contiguousRhs, contiguousOutput);
     return success();
   }
 };
@@ -546,8 +551,9 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
       return failure();
 
     Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
+    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
-    replaceOpWithNewBufferizedOp<OpTy>(rewriter, op, outputBufferOpt->getType(), contiguousInput, *outputBufferOpt);
+    replaceOpWithNewBufferizedOp<OpTy>(rewriter, op, contiguousOutput.getType(), contiguousInput, contiguousOutput);
     return success();
   }
 };

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -240,10 +240,10 @@ void SpatCompute::print(OpAsmPrinter& printer) {
   printBoundValueList(printer, weightArgs, getWeights(), ListDelimiter::Square);
   printer << " ";
   printBoundValueList(printer, inputArgs, getInputs(), ListDelimiter::Paren);
-  printer << " crossbarWeights " << collectDistinctCrossbarWeights(getOperation()).size();
 
   if (auto coreIdAttr = (*this)->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
     printer << " coreId " << coreIdAttr.getInt();
+  printer << " crossbarWeights " << collectDistinctCrossbarWeights(getOperation()).size();
 
   printer.printOptionalAttrDict((*this)->getAttrs(),
                                 {getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdAttrName});
@@ -276,13 +276,13 @@ ParseResult SpatCompute::parse(OpAsmParser& parser, OperationState& result) {
   if (parseBoundValueList(parser, ListDelimiter::Paren, inputArgs, inputs))
     return failure();
 
-  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
-  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
-    return failure();
-
   bool hasCoreId = parseOptionalKeywordAlias(parser, "coreId", "core_id");
   if (hasCoreId && parser.parseInteger(coreId))
     return failure();
+
+  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
+  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
+    return failure();
   (void) crossbarWeightCount;
 
   if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
@@ -365,13 +365,14 @@ void SpatComputeBatch::print(OpAsmPrinter& printer) {
   printBoundValueList(printer, weightArgs, getWeights(), ListDelimiter::Square);
   printer << " ";
   printBoundValueList(printer, inputArgs, getInputs(), ListDelimiter::Paren);
-  printer << " crossbarWeights " << collectDistinctCrossbarWeights(getOperation()).size();
 
   if (getNumResults() != 0) {
     printer << " shared_outs";
     printBlockArgumentList(printer, outputArgs);
   }
 
+  printer << " crossbarWeights " << getComputeInstanceCrossbarUsage({getOperation(), 0, getLaneCount()}).size();
+
   if (auto coreIdsAttr = (*this)->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
     printer << " coreIds ";
     printCompressedIntegerList(printer, coreIdsAttr.asArrayRef());
@@ -423,13 +424,13 @@ ParseResult SpatComputeBatch::parse(OpAsmParser& parser, OperationState& result)
     if (parseBlockArgumentList(parser, outputArgs))
       return failure();
 
-  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
-  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
-    return failure();
-
   bool hasCoreIds = parseOptionalKeywordAlias(parser, "coreIds", "core_ids");
   if (hasCoreIds && parseCompressedIntegerList(parser, coreIds))
     return failure();
+
+  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
+  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
+    return failure();
   (void) crossbarWeightCount;
 
   if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()