Skip to content

Commit

Permalink
extend to all doubly strided ops
Browse files Browse the repository at this point in the history
  • Loading branch information
@newling
newling committed Aug 16, 2024
1 parent 1de7f26 commit 5f1696c
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 107 deletions.
29 changes: 29 additions & 0 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -465,6 +465,35 @@ void NpuDmaCpyNdOp::getCanonicalizationPatterns(RewritePatternSet &results,
context);
}

template <typename T>
struct DmaCpyNdBaseOpReplacementBuilder {
static void replace(T dmaOp, PatternRewriter &rewriter,
ArrayRef<OpFoldResult> tgtMixedOffsets,
ArrayRef<OpFoldResult> tgtMixedSizes,
ArrayRef<OpFoldResult> tgtMixedStrides,
ArrayRef<OpFoldResult> srcMixedOffsets,
ArrayRef<OpFoldResult> srcMixedSizes,
ArrayRef<OpFoldResult> srcMixedStrides) {
rewriter.replaceOpWithNewOp<T>(dmaOp, dmaOp.getTarget(), tgtMixedOffsets,
tgtMixedSizes, tgtMixedStrides,
dmaOp.getSource(), srcMixedOffsets,
srcMixedSizes, srcMixedStrides);
}
};

void DmaCpyNdOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<DoublyStridedFolder<
DmaCpyNdOp, DmaCpyNdBaseOpReplacementBuilder<DmaCpyNdOp>>>(context);
}

void CircularDmaCpyNdOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<DoublyStridedFolder<
CircularDmaCpyNdOp,
DmaCpyNdBaseOpReplacementBuilder<CircularDmaCpyNdOp>>>(context);
}

LogicalResult LogicalObjectFifoFromMemrefOp::canonicalize(
LogicalObjectFifoFromMemrefOp logicalObjectFifo,
PatternRewriter &rewriter) {
Expand Down
11 changes: 2 additions & 9 deletions compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -380,7 +380,6 @@ def AMDAIE_NpuDmaCpyNdOp: AMDAIE_Op<"npu.dma_cpy_nd",

}];

// Ensure that dimensions of offsets/sizes/strides that can be static, are.
let hasCanonicalizer = 1;

}
Expand Down Expand Up @@ -815,10 +814,7 @@ def AMDAIE_DmaCpyNdOp: AMDAIE_DmaCpyNdBaseOp<"dma_cpy_nd", []> {
}];

let hasVerifier = 0;

// Ensure that dimensions of offsets/sizes/strides that can be static, are.
// TODO(newling)
// let hasCanonicalizer = 1;
let hasCanonicalizer = 1;
}

def AMDAIE_CircularDmaCpyNdOp: AMDAIE_DmaCpyNdBaseOp<"circular_dma_cpy_nd", [Pure]> {
Expand All @@ -843,10 +839,7 @@ def AMDAIE_CircularDmaCpyNdOp: AMDAIE_DmaCpyNdBaseOp<"circular_dma_cpy_nd", [Pur
}];

let hasVerifier = 0;

// Ensure that dimensions of offsets/sizes/strides that can be static, are.
// TODO(newling)
// let hasCanonicalizer = 1;
let hasCanonicalizer = 1;
}

def AMDAIE_ReferenceToOp: AMDAIE_Op<"reference_to", [SameOperandsAndResultType]> {
Expand Down
116 changes: 18 additions & 98 deletions ...r/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/canonicalize_doubly_strided_op.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,7 @@
// CHECK-LABEL: func.func @circular_dma_cpy_nd_source_target
// CHECK-SAME: %[[ARG0:.+]]: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
// CHECK-SAME: %[[ARG1:.+]]: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.circular_dma_cpy_nd(%[[ARG0]][%[[C0]]] [%[[C128]]] [%[[C1]]], %[[ARG1]][%[[C0]]] [%[[C64]]] [%[[C1]]])
// CHECK: amdaie.circular_dma_cpy_nd(%[[ARG0]][0] [128] [1], %[[ARG1]][0] [64] [1])

// FOLD-SINGLE-DIMS-LABEL: func.func @circular_dma_cpy_nd_source_target
// FOLD-SINGLE-DIMS-SAME: %[[ARG0:.+]]: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
Expand All @@ -25,11 +21,7 @@ func.func @circular_dma_cpy_nd_source_target(%arg0: !amdaie.logicalobjectfifo<me
// -----

// CHECK-LABEL: func.func @circular_dma_cpy_nd_linear_implicit
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C0]]] [%[[C64]]] [%[[C1]]])
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[0] [128] [1], %{{.+}}[0] [64] [1])
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[] [] [])
func.func @circular_dma_cpy_nd_linear_implicit(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
%0 = amdaie.circular_dma_cpy_nd(%arg0[0, 0, 0, 0] [1, 1, 8, 16] [128, 128, 16, 1], %arg1[0, 0, 0, 0] [1, 4, 2, 8] [64, 16, 8, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
Expand All @@ -40,20 +32,8 @@ func.func @circular_dma_cpy_nd_linear_implicit(%arg0: !amdaie.logicalobjectfifo<
// -----

// CHECK-LABEL: func.func @circular_dma_cpy_nd_linear
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C0]], %[[C0]]] [%[[C16]], %[[C8]]] [%[[C16]], %[[C1]]], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C64]], %[[C16]], %[[C128]]] [%[[C128]], %[[C16]], %[[C1]]])
// FOLD-SINGLE-DIMS-DAG: %[[C0:.+]] = arith.constant 0 : index
// FOLD-SINGLE-DIMS-DAG: %[[C1:.+]] = arith.constant 1 : index
// FOLD-SINGLE-DIMS-DAG: %[[C8:.+]] = arith.constant 8 : index
// FOLD-SINGLE-DIMS-DAG: %[[C16:.+]] = arith.constant 16 : index
// FOLD-SINGLE-DIMS-DAG: %[[C64:.+]] = arith.constant 64 : index
// FOLD-SINGLE-DIMS-DAG: %[[C128:.+]] = arith.constant 128 : index
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C0]], %[[C0]]] [%[[C16]], %[[C8]]] [%[[C16]], %[[C1]]], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C64]], %[[C16]], %[[C128]]] [%[[C128]], %[[C16]], %[[C1]]])
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[0, 0] [16, 8] [16, 1], %{{.+}}[0, 0, 0] [64, 16, 128] [128, 16, 1])
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[0, 0] [16, 8] [16, 1], %{{.+}}[0, 0, 0] [64, 16, 128] [128, 16, 1])
func.func @circular_dma_cpy_nd_linear(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
%c16 = arith.constant 16 : index
%0 = amdaie.circular_dma_cpy_nd(%arg0[0, 0, 0, 0] [1, 2, 8, 8] [256, 128, %c16, 1], %arg1[0, 0, 0, 0] [64, 16, 8, %c16] [128, %c16, %c16, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
Expand All @@ -75,19 +55,8 @@ func.func @circular_dma_cpy_nd_no_linear(%arg0: !amdaie.logicalobjectfifo<memref
// -----

// CHECK-LABEL: func.func @circular_dma_cpy_nd_unit
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C2]], %[[C8]], %[[C8]]] [%[[C8]], %[[C16]], %[[C1]]])
// FOLD-SINGLE-DIMS-DAG: %[[C0:.+]] = arith.constant 0 : index
// FOLD-SINGLE-DIMS-DAG: %[[C1:.+]] = arith.constant 1 : index
// FOLD-SINGLE-DIMS-DAG: %[[C2:.+]] = arith.constant 2 : index
// FOLD-SINGLE-DIMS-DAG: %[[C8:.+]] = arith.constant 8 : index
// FOLD-SINGLE-DIMS-DAG: %[[C16:.+]] = arith.constant 16 : index
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C2]], %[[C8]], %[[C8]]] [%[[C8]], %[[C16]], %[[C1]]])
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[0] [128] [1], %{{.+}}[0, 0, 0] [2, 8, 8] [8, 16, 1])
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[0, 0, 0] [2, 8, 8] [8, 16, 1])
func.func @circular_dma_cpy_nd_unit(%arg0: !amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>) {
%0 = amdaie.circular_dma_cpy_nd(%arg0[0, 0, 0, 0, 0, 0] [1, 1, 2, 2, 4, 8] [128, 128, 64, 32, 8, 1], %arg1[0, 0, 0, 0, 0, 0] [1, 1, 2, 2, 4, 8] [128, 128, 8, 64, 16, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>)
"iree.keep"(%0) : (index) -> ()
Expand All @@ -97,10 +66,7 @@ func.func @circular_dma_cpy_nd_unit(%arg0: !amdaie.logicalobjectfifo<memref<1x1x
// -----

// CHECK-LABEL: func.func @circular_dma_cpy_nd_unit_between_linear
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]])
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[0] [128] [1], %{{.+}}[0] [128] [1])
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[] [] [])
func.func @circular_dma_cpy_nd_unit_between_linear(%arg0: !amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>) {
%0 = amdaie.circular_dma_cpy_nd(%arg0[0, 0, 0, 0, 0, 0] [1, 2, 2, 4, 1, 8] [128, 64, 32, 8, 8, 1], %arg1[0, 0, 0, 0, 0, 0] [2, 2, 1, 4, 8, 1] [64, 32, 32, 8, 1, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>)
Expand All @@ -122,14 +88,8 @@ func.func @circular_dma_cpy_nd_non_zero_offset(%arg0: !amdaie.logicalobjectfifo<
// -----

// CHECK-LABEL: func.func @circular_dma_cpy_nd_partial_non_zero_offset
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C1]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C1]]] [%[[C64]]] [%[[C1]]])
// FOLD-SINGLE-DIMS-DAG: %[[C1:.+]] = arith.constant 1 : index
// FOLD-SINGLE-DIMS-DAG: %[[C64:.+]] = arith.constant 64 : index
// FOLD-SINGLE-DIMS-DAG: %[[C128:.+]] = arith.constant 128 : index
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[%[[C1]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C1]]] [%[[C64]]] [%[[C1]]])
// CHECK: amdaie.circular_dma_cpy_nd(%{{.+}}[1] [128] [1], %{{.+}}[1] [64] [1])
// FOLD-SINGLE-DIMS: amdaie.circular_dma_cpy_nd(%{{.+}}[1] [128] [1], %{{.+}}[1] [64] [1])
func.func @circular_dma_cpy_nd_partial_non_zero_offset(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
%0 = amdaie.circular_dma_cpy_nd(%arg0[0, 0, 0, 1] [1, 1, 8, 16] [128, 128, 16, 1], %arg1[0, 0, 0, 1] [1, 4, 2, 8] [64, 16, 8, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
"iree.keep"(%0) : (index) -> ()
Expand All @@ -143,11 +103,7 @@ func.func @circular_dma_cpy_nd_partial_non_zero_offset(%arg0: !amdaie.logicalobj
// CHECK-LABEL: func.func @dma_cpy_nd_source_target
// CHECK-SAME: %[[ARG0:.+]]: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
// CHECK-SAME: %[[ARG1:.+]]: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.dma_cpy_nd(%[[ARG0]][%[[C0]]] [%[[C128]]] [%[[C1]]], %[[ARG1]][%[[C0]]] [%[[C64]]] [%[[C1]]])
// CHECK: amdaie.dma_cpy_nd(%[[ARG0]][0] [128] [1], %[[ARG1]][0] [64] [1])

// FOLD-SINGLE-DIMS-LABEL: func.func @dma_cpy_nd_source_target
// FOLD-SINGLE-DIMS-SAME: %[[ARG0:.+]]: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>
Expand All @@ -162,11 +118,7 @@ func.func @dma_cpy_nd_source_target(%arg0: !amdaie.logicalobjectfifo<memref<1x1x
// -----

// CHECK-LABEL: func.func @dma_cpy_nd_linear_implicit
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C0]]] [%[[C64]]] [%[[C1]]])
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[0] [128] [1], %{{.+}}[0] [64] [1])
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[] [] [])
func.func @dma_cpy_nd_linear_implicit(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
%0 = amdaie.dma_cpy_nd(%arg0[0, 0, 0, 0] [1, 1, 8, 16] [128, 128, 16, 1], %arg1[0, 0, 0, 0] [1, 4, 2, 8] [64, 16, 8, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
Expand All @@ -177,20 +129,8 @@ func.func @dma_cpy_nd_linear_implicit(%arg0: !amdaie.logicalobjectfifo<memref<1x
// -----

// CHECK-LABEL: func.func @dma_cpy_nd_linear
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[%[[C0]], %[[C0]]] [%[[C16]], %[[C8]]] [%[[C16]], %[[C1]]], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C64]], %[[C16]], %[[C128]]] [%[[C128]], %[[C16]], %[[C1]]])
// FOLD-SINGLE-DIMS-DAG: %[[C0:.+]] = arith.constant 0 : index
// FOLD-SINGLE-DIMS-DAG: %[[C1:.+]] = arith.constant 1 : index
// FOLD-SINGLE-DIMS-DAG: %[[C8:.+]] = arith.constant 8 : index
// FOLD-SINGLE-DIMS-DAG: %[[C16:.+]] = arith.constant 16 : index
// FOLD-SINGLE-DIMS-DAG: %[[C64:.+]] = arith.constant 64 : index
// FOLD-SINGLE-DIMS-DAG: %[[C128:.+]] = arith.constant 128 : index
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[%[[C0]], %[[C0]]] [%[[C16]], %[[C8]]] [%[[C16]], %[[C1]]], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C64]], %[[C16]], %[[C128]]] [%[[C128]], %[[C16]], %[[C1]]])
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[0, 0] [16, 8] [16, 1], %{{.+}}[0, 0, 0] [64, 16, 128] [128, 16, 1])
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[0, 0] [16, 8] [16, 1], %{{.+}}[0, 0, 0] [64, 16, 128] [128, 16, 1])
func.func @dma_cpy_nd_linear(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
%c16 = arith.constant 16 : index
%0 = amdaie.dma_cpy_nd(%arg0[0, 0, 0, 0] [1, 2, 8, 8] [256, 128, %c16, 1], %arg1[0, 0, 0, 0] [64, 16, 8, %c16] [128, %c16, %c16, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
Expand All @@ -212,19 +152,8 @@ func.func @dma_cpy_nd_no_linear(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16
// -----

// CHECK-LABEL: func.func @dma_cpy_nd_unit
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C2]], %[[C8]], %[[C8]]] [%[[C8]], %[[C16]], %[[C1]]])
// FOLD-SINGLE-DIMS-DAG: %[[C0:.+]] = arith.constant 0 : index
// FOLD-SINGLE-DIMS-DAG: %[[C1:.+]] = arith.constant 1 : index
// FOLD-SINGLE-DIMS-DAG: %[[C2:.+]] = arith.constant 2 : index
// FOLD-SINGLE-DIMS-DAG: %[[C8:.+]] = arith.constant 8 : index
// FOLD-SINGLE-DIMS-DAG: %[[C16:.+]] = arith.constant 16 : index
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[%[[C0]], %[[C0]], %[[C0]]] [%[[C2]], %[[C8]], %[[C8]]] [%[[C8]], %[[C16]], %[[C1]]])
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[0] [128] [1], %{{.+}}[0, 0, 0] [2, 8, 8] [8, 16, 1])
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[0, 0, 0] [2, 8, 8] [8, 16, 1])
func.func @dma_cpy_nd_unit(%arg0: !amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>) {
%0 = amdaie.dma_cpy_nd(%arg0[0, 0, 0, 0, 0, 0] [1, 1, 2, 2, 4, 8] [128, 128, 64, 32, 8, 1], %arg1[0, 0, 0, 0, 0, 0] [1, 1, 2, 2, 4, 8] [128, 128, 8, 64, 16, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>)
"iree.keep"(%0) : (index) -> ()
Expand All @@ -234,10 +163,7 @@ func.func @dma_cpy_nd_unit(%arg0: !amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi
// -----

// CHECK-LABEL: func.func @dma_cpy_nd_unit_between_linear
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C0]]] [%[[C128]]] [%[[C1]]])
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[0] [128] [1], %{{.+}}[0] [128] [1])
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[] [] [], %{{.+}}[] [] [])
func.func @dma_cpy_nd_unit_between_linear(%arg0: !amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>) {
%0 = amdaie.dma_cpy_nd(%arg0[0, 0, 0, 0, 0, 0] [2, 2, 1, 1, 4, 8] [64, 32, 32, 32, 8, 1], %arg1[0, 0, 0, 0, 0, 0] [2, 1, 2, 1, 4, 8] [64, 64, 32, 32, 8, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x2x2x4x8xi32, 1>>, !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>)
Expand All @@ -259,14 +185,8 @@ func.func @dma_cpy_nd_non_zero_offset(%arg0: !amdaie.logicalobjectfifo<memref<1x
// -----

// CHECK-LABEL: func.func @dma_cpy_nd_partial_non_zero_offset
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[%[[C1]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C1]]] [%[[C64]]] [%[[C1]]])
// FOLD-SINGLE-DIMS-DAG: %[[C1:.+]] = arith.constant 1 : index
// FOLD-SINGLE-DIMS-DAG: %[[C64:.+]] = arith.constant 64 : index
// FOLD-SINGLE-DIMS-DAG: %[[C128:.+]] = arith.constant 128 : index
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[%[[C1]]] [%[[C128]]] [%[[C1]]], %{{.+}}[%[[C1]]] [%[[C64]]] [%[[C1]]])
// CHECK: amdaie.dma_cpy_nd(%{{.+}}[1] [128] [1], %{{.+}}[1] [64] [1])
// FOLD-SINGLE-DIMS: amdaie.dma_cpy_nd(%{{.+}}[1] [128] [1], %{{.+}}[1] [64] [1])
func.func @dma_cpy_nd_partial_non_zero_offset(%arg0: !amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, %arg1: !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>) {
%0 = amdaie.dma_cpy_nd(%arg0[0, 0, 0, 1] [1, 1, 8, 16] [128, 128, 16, 1], %arg1[0, 0, 0, 1] [1, 4, 2, 8] [64, 16, 8, 1]) : (!amdaie.logicalobjectfifo<memref<1x1x8x16xi32, 1>>, !amdaie.logicalobjectfifo<memref<8x16xi32, 1>>)
"iree.keep"(%0) : (index) -> ()
Expand Down

0 comments on commit 5f1696c

Please sign in to comment.