29
29
#include " mlir/Support/LLVM.h"
30
30
#include " llvm/ADT/TypeSwitch.h"
31
31
#include " llvm/Support/Debug.h"
32
+ #include " llvm/Support/DebugLog.h"
32
33
#include " llvm/Support/InterleavedRange.h"
33
34
#include " llvm/Support/raw_ostream.h"
34
35
#include < utility>
38
39
using namespace mlir ;
39
40
using namespace mlir ::linalg;
40
41
41
- #define DBGS () (llvm::dbgs() << " [" DEBUG_TYPE << " ]: " )
42
- #define DBGSNL () (llvm::dbgs() << " \n " )
43
-
44
42
// ===----------------------------------------------------------------------===//
45
43
// Transformations exposed as functional-style API calls.
46
44
// ===----------------------------------------------------------------------===//
@@ -276,23 +274,18 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
276
274
tensor::PadOp::create (rewriter, loc, collapsed, packOp.getSource (), lows,
277
275
highs, paddingValue, /* nofold=*/ false );
278
276
279
- LLVM_DEBUG (
280
- DBGSNL (); DBGSNL ();
281
- DBGS () << " insertPositions: "
282
- << llvm::interleaved (packingMetadata.insertPositions );
283
- DBGSNL (); DBGS () << " outerPositions: "
284
- << llvm::interleaved (packingMetadata.outerPositions );
285
- DBGSNL (); DBGS () << " packedShape: "
286
- << llvm::interleaved (packedTensorType.getShape ());
287
- DBGSNL (); DBGS () << " packedToStripMinedShapePerm: "
288
- << llvm::interleaved (packedToStripMinedShapePerm);
289
- DBGSNL ();
290
- DBGS () << " reassociations: "
291
- << llvm::interleaved (llvm::map_range (
292
- packingMetadata.reassociations , stringifyReassocIndices));
293
- DBGSNL ();
294
- DBGS () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
295
- DBGSNL (); DBGS () << " collapsed type: " << collapsed; DBGSNL (););
277
+ LDBG () << " insertPositions: "
278
+ << llvm::interleaved (packingMetadata.insertPositions );
279
+ LDBG () << " outerPositions: "
280
+ << llvm::interleaved (packingMetadata.outerPositions );
281
+ LDBG () << " packedShape: " << llvm::interleaved (packedTensorType.getShape ());
282
+ LDBG () << " packedToStripMinedShapePerm: "
283
+ << llvm::interleaved (packedToStripMinedShapePerm);
284
+ LDBG () << " reassociations: "
285
+ << llvm::interleaved (llvm::map_range (packingMetadata.reassociations ,
286
+ stringifyReassocIndices));
287
+ LDBG () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
288
+ LDBG () << " collapsed type: " << collapsed;
296
289
297
290
if (lowerPadLikeWithInsertSlice && packOp.isLikePad ()) {
298
291
// Pack ops which operate as simple pads may not produce legal
@@ -317,7 +310,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
317
310
rewriter, loc, /* source=*/ padOp, /* dest=*/ packOp.getDest (),
318
311
/* offsets=*/ zeros, sizes, /* strides=*/ ones);
319
312
320
- LLVM_DEBUG ( DBGS ( ) << " insert_slice op: " << insertSliceOp; DBGSNL ();) ;
313
+ LDBG ( ) << " insert_slice op: " << insertSliceOp;
321
314
322
315
rewriter.replaceOp (packOp, insertSliceOp->getResults ());
323
316
@@ -339,10 +332,9 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
339
332
auto transposeOp = linalg::TransposeOp::create (
340
333
rewriter, loc, reshapeOp.getResult (), packOp.getDest (), transpPerm);
341
334
342
- LLVM_DEBUG (DBGSNL (); DBGSNL (); DBGSNL ();
343
- DBGS () << " reshape op: " << reshapeOp; DBGSNL ();
344
- DBGS () << " transpPerm: " << llvm::interleaved (transpPerm);
345
- DBGSNL (); DBGS () << " transpose op: " << transposeOp; DBGSNL (););
335
+ LDBG () << " reshape op: " << reshapeOp;
336
+ LDBG () << " transpPerm: " << llvm::interleaved (transpPerm);
337
+ LDBG () << " transpose op: " << transposeOp;
346
338
347
339
// 7. Replace packOp by transposeOp.
348
340
rewriter.replaceOp (packOp, transposeOp->getResults ());
@@ -410,21 +402,16 @@ linalg::lowerUnPack(RewriterBase &rewriter, linalg::UnPackOp unPackOp,
410
402
linalg::TransposeOp::create (rewriter, loc, unPackOp.getSource (), emptyOp,
411
403
packedToStripMinedShapePerm);
412
404
413
- LLVM_DEBUG (
414
- DBGSNL (); DBGSNL ();
415
- DBGS () << " insertPositions: "
416
- << llvm::interleaved (packingMetadata.insertPositions );
417
- DBGSNL (); DBGS () << " packedShape: "
418
- << llvm::interleaved (packedTensorType.getShape ());
419
- DBGSNL (); DBGS () << " packedToStripMinedShapePerm: "
420
- << llvm::interleaved (packedToStripMinedShapePerm);
421
- DBGSNL ();
422
- DBGS () << " reassociations: "
423
- << llvm::interleaved (llvm::map_range (
424
- packingMetadata.reassociations , stringifyReassocIndices));
425
- DBGSNL ();
426
- DBGS () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
427
- DBGSNL (); DBGS () << " collapsed type: " << collapsedType; DBGSNL (););
405
+ LDBG () << " insertPositions: "
406
+ << llvm::interleaved (packingMetadata.insertPositions );
407
+ LDBG () << " packedShape: " << llvm::interleaved (packedTensorType.getShape ());
408
+ LDBG () << " packedToStripMinedShapePerm: "
409
+ << llvm::interleaved (packedToStripMinedShapePerm);
410
+ LDBG () << " reassociations: "
411
+ << llvm::interleaved (llvm::map_range (packingMetadata.reassociations ,
412
+ stringifyReassocIndices));
413
+ LDBG () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
414
+ LDBG () << " collapsed type: " << collapsedType;
428
415
429
416
// 4. Collapse from the stripMinedShape to the padded result.
430
417
auto reshapeOp = tensor::CollapseShapeOp::create (
@@ -486,10 +473,9 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
486
473
SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray ();
487
474
SmallVector<utils::IteratorType> iteratorTypes =
488
475
linalgOp.getIteratorTypesArray ();
489
- LLVM_DEBUG (DBGS () << " Start packing: " << linalgOp << " \n "
490
- << " maps: " << llvm::interleaved (indexingMaps) << " \n "
491
- << " iterators: " << llvm::interleaved (iteratorTypes)
492
- << " \n " );
476
+ LDBG () << " Start packing: " << linalgOp;
477
+ LDBG () << " maps: " << llvm::interleaved (indexingMaps);
478
+ LDBG () << " iterators: " << llvm::interleaved (iteratorTypes);
493
479
494
480
SmallVector<linalg::PackOp> packOps;
495
481
SmallVector<linalg::UnPackOp> unPackOps;
@@ -511,14 +497,11 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
511
497
packedOperandsDims.packedDimForEachOperand = *maybePackedDimForEachOperand;
512
498
listOfPackedOperandsDim.pushBack (std::move (packedOperandsDims));
513
499
514
- LLVM_DEBUG (
515
- DBGS () << " ++++ After pack size #" << i << " : " << packedSizes[i]
516
- << " \n "
517
- << " maps: " << llvm::interleaved (indexingMaps) << " \n "
518
- << " iterators: " << llvm::interleaved (iteratorTypes) << " \n "
519
- << " packedDimForEachOperand: "
520
- << llvm::interleaved (packedOperandsDims.packedDimForEachOperand )
521
- << " \n " );
500
+ LDBG () << " ++++ After pack size #" << i << " : " << packedSizes[i];
501
+ LDBG () << " maps: " << llvm::interleaved (indexingMaps);
502
+ LDBG () << " iterators: " << llvm::interleaved (iteratorTypes);
503
+ LDBG () << " packedDimForEachOperand: "
504
+ << llvm::interleaved (packedOperandsDims.packedDimForEachOperand );
522
505
}
523
506
524
507
// Step 2. Propagate packing to all LinalgOp operands.
@@ -534,10 +517,9 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
534
517
listOfPackedOperandsDim.extractPackedDimsForOperand (pos);
535
518
SmallVector<OpFoldResult> innerPackSizes =
536
519
listOfPackedOperandsDim.extractPackSizesForOperand (pos);
537
- LLVM_DEBUG (DBGS () << " operand: " << operand << " \n "
538
- << " innerPos: " << llvm::interleaved (innerPos) << " \n "
539
- << " innerPackSizes: "
540
- << llvm::interleaved (innerPackSizes) << " \n " );
520
+ LDBG () << " operand: " << operand;
521
+ LDBG () << " innerPos: " << llvm::interleaved (innerPos);
522
+ LDBG () << " innerPackSizes: " << llvm::interleaved (innerPackSizes);
541
523
if (innerPackSizes.empty ()) {
542
524
inputsAndInits.push_back (operand);
543
525
continue ;
@@ -776,8 +758,8 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
776
758
777
759
int64_t numLoops = linalgOp.getNumLoops ();
778
760
if (numLoops <= 2 ) {
779
- LLVM_DEBUG ( DBGS ( ) << " need 3+ loops to find a matmul to pack, got "
780
- << numLoops << " \n in : " << linalgOp << " \n " ) ;
761
+ LDBG ( ) << " need 3+ loops to find a matmul to pack, got " << numLoops
762
+ << " in : " << linalgOp;
781
763
return rewriter.notifyMatchFailure (
782
764
linalgOp, " need 3+ loops to find a matmul to pack" );
783
765
}
@@ -801,8 +783,7 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
801
783
FailureOr<ContractionDimensions> maybeDimensions =
802
784
inferContractionDims (linalgOp);
803
785
if (failed (maybeDimensions)) {
804
- LLVM_DEBUG (DBGS () << " couldn't infer matmul iterators in: " << linalgOp
805
- << " \n " );
786
+ LDBG () << " couldn't infer matmul iterators in: " << linalgOp;
806
787
return rewriter.notifyMatchFailure (linalgOp,
807
788
" couldn't infer matmul iterators" );
808
789
}
@@ -814,10 +795,8 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
814
795
// to plug a heuristic.
815
796
int64_t mPos = maybeDimensions->m .back (), nPos = maybeDimensions->n .back (),
816
797
kPos = maybeDimensions->k .back ();
817
- LLVM_DEBUG (DBGSNL (); DBGSNL (); DBGSNL ();
818
- DBGS () << " Start packing generic op greedily with (m@" << mPos
819
- << " , n@" << nPos << " , k@" << kPos << " ): " << linalgOp
820
- << " \n " ;);
798
+ LDBG () << " Start packing generic op greedily with (m@" << mPos << " , n@"
799
+ << nPos << " , k@" << kPos << " ): " << linalgOp;
821
800
822
801
// 2.a. Rewrite as a generic.
823
802
auto genericOp = dyn_cast<GenericOp>(linalgOp.getOperation ());
@@ -833,14 +812,14 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
833
812
// not change the indexings of any operand.
834
813
SmallVector<int64_t > permutation =
835
814
computePermutationVector (numLoops, {mPos , nPos, kPos }, mmnnkkPos);
836
- LLVM_DEBUG ( DBGS ( ) << " perm: " << llvm::interleaved (permutation) << " \n " );
815
+ LDBG ( ) << " perm: " << llvm::interleaved (permutation);
837
816
// Sign .. unsigned pollution.
838
817
SmallVector<unsigned > unsignedPerm (permutation.begin (), permutation.end ());
839
818
FailureOr<GenericOp> interchangeResult =
840
819
interchangeGenericOp (rewriter, genericOp, unsignedPerm);
841
820
assert (succeeded (interchangeResult) && " unexpected failure interchanging op" );
842
821
genericOp = *interchangeResult;
843
- LLVM_DEBUG ( DBGS ( ) << " Generalized Op to pack: " << genericOp << " \n " ;) ;
822
+ LDBG ( ) << " Generalized Op to pack: " << genericOp;
844
823
845
824
// At this point, the op iterators are normalized to {leading, k, m, n}.
846
825
// The layouts induced by packing will always be:
@@ -862,12 +841,11 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
862
841
863
842
// Add leading zeros to match numLoops, we only pack the last 3 dimensions
864
843
// post interchange.
865
- LLVM_DEBUG (DBGS () << " paddedSizesNextMultipleOf: "
866
- << llvm::interleaved (paddedSizesNextMultipleOf) << " \n "
867
- << " loopRanges: "
868
- << llvm::interleaved (llvm::map_range (
869
- loopRanges, [](Range r) { return r.size ; }))
870
- << " \n " );
844
+ LDBG () << " paddedSizesNextMultipleOf: "
845
+ << llvm::interleaved (paddedSizesNextMultipleOf);
846
+ LDBG () << " loopRanges: "
847
+ << llvm::interleaved (
848
+ llvm::map_range (loopRanges, [](Range r) { return r.size ; }));
871
849
SmallVector<OpFoldResult> adjustedPackedSizes (numLoops - packedSizes.size (),
872
850
rewriter.getIndexAttr (0 ));
873
851
for (int64_t i = 0 , e = numPackedDims; i < e; ++i) {
@@ -883,8 +861,7 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
883
861
{loopRanges[adjustedPackedSizes.size ()].size ,
884
862
rewriter.getIndexAttr (paddedSizesNextMultipleOf[i])}));
885
863
}
886
- LLVM_DEBUG (DBGS () << " adjustedPackedSizes: "
887
- << llvm::interleaved (adjustedPackedSizes) << " \n " );
864
+ LDBG () << " adjustedPackedSizes: " << llvm::interleaved (adjustedPackedSizes);
888
865
889
866
// TODO: If we wanted to give the genericOp a name after packing, after
890
867
// calling `pack` would be a good time. One would still need to check that
@@ -1214,9 +1191,8 @@ LogicalResult DecomposeOuterUnitDimsPackOpPattern::matchAndRewrite(
1214
1191
}
1215
1192
srcPermForTranspose.append (innerDimPos.begin (), innerDimPos.end ());
1216
1193
1217
- LLVM_DEBUG (DBGS () << " Pack permutation: " << packOp << " \n "
1218
- << " perm: " << llvm::interleaved (srcPermForTranspose)
1219
- << " \n " );
1194
+ LDBG () << " Pack permutation: " << packOp;
1195
+ LDBG () << " perm: " << llvm::interleaved (srcPermForTranspose);
1220
1196
1221
1197
// 2.1 Create tensor.empty (init value for TransposeOp)
1222
1198
SmallVector<OpFoldResult> transShapeForEmptyOp (srcRank - numTiles,
0 commit comments