From 46b755d4056f6bab9aee9a5d1dd37217d454e172 Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Mon, 24 Jun 2019 09:00:06 -0700 Subject: [PATCH] Use linalg.view_slice in tiling and fusion This CL makes use of view_slice in tiling and fusion. Using a higher level IR element greatly simplifies the IR produced during tiling and fusion. Lowering to LLVM is updated to first translate view_slice into a sequence of dim, range and cmpi. This level will also be useful when lowering to affine. PiperOrigin-RevId: 254767814 --- mlir/include/mlir/Linalg/IR/LinalgOps.h | 2 + mlir/include/mlir/Linalg/IR/LinalgOps.td | 18 ++ mlir/lib/Linalg/IR/LinalgOps.cpp | 7 + mlir/lib/Linalg/Transforms/Fusion.cpp | 45 +++-- mlir/lib/Linalg/Transforms/Tiling.cpp | 44 ++--- mlir/test/Linalg/roundtrip.mlir | 2 +- mlir/test/Linalg/tile.mlir | 280 ++++++++----------------------- 7 files changed, 126 insertions(+), 272 deletions(-) diff --git a/mlir/include/mlir/Linalg/IR/LinalgOps.h b/mlir/include/mlir/Linalg/IR/LinalgOps.h index 5574d41..ff04c36 100644 --- a/mlir/include/mlir/Linalg/IR/LinalgOps.h +++ b/mlir/include/mlir/Linalg/IR/LinalgOps.h @@ -356,6 +356,8 @@ public: #define GET_OP_CLASSES #include "mlir/Linalg/IR/LinalgLibraryOps.h.inc" +llvm::raw_ostream &operator<<(llvm::raw_ostream &os, SubViewOp::Range &range); + /// Returns the list of maps that map loops to operands of a Linalg op. /// The i-th affine map identifies loop indices to subscripts that are used when /// accessing the i-th operand. diff --git a/mlir/include/mlir/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Linalg/IR/LinalgOps.td index 4b582aa..051d16e 100644 --- a/mlir/include/mlir/Linalg/IR/LinalgOps.td +++ b/mlir/include/mlir/Linalg/IR/LinalgOps.td @@ -153,6 +153,15 @@ def SubViewOp : Linalg_Op<"subview", [NoSideEffect]>, // TODO(ntv) evolve syntax towards: // linalg.subview %0[%1:%2:%3][%4:%5:%6] : view + let builders = [OpBuilder< + "Builder *builder, OperationState *result, Value *view, " + "ArrayRef ranges", + [{ + result->addOperands(view); + result->addOperands(ranges); + result->types.push_back(view->getType()); + }]>]; + let verifier = [{ auto numRanges = (getNumOperands() - 1) / 3; if (getNumOperands() != 3 * numRanges + 1 || @@ -178,6 +187,15 @@ def SubViewOp : Linalg_Op<"subview", [NoSideEffect]>, res.push_back(getRange(i)); return res; } + // This requires `SubViewOp` to be declared, in the future it should be + // folded into the builders. + static void build(Builder *builder, OperationState *result, Value *view, + ArrayRef ranges) { + result->addOperands(view); + for (auto r : ranges) + result->addOperands({r.min, r.max, r.step}); + result->types.push_back(view->getType()); + } }]; } diff --git a/mlir/lib/Linalg/IR/LinalgOps.cpp b/mlir/lib/Linalg/IR/LinalgOps.cpp index ded07f1..0c7c019 100644 --- a/mlir/lib/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Linalg/IR/LinalgOps.cpp @@ -814,6 +814,12 @@ static LogicalResult verify(ConvOp op) { return success(); } +llvm::raw_ostream &mlir::linalg::operator<<(llvm::raw_ostream &os, + SubViewOp::Range &range) { + return os << "range " << *range.min << ":" << *range.max << ":" + << *range.step; +} + namespace mlir { namespace linalg { @@ -825,6 +831,7 @@ namespace linalg { } // namespace linalg } // namespace mlir + static AffineMap extractOrIdentityMap(llvm::Optional maybeMap, unsigned rank, MLIRContext *context) { if (maybeMap) diff --git a/mlir/lib/Linalg/Transforms/Fusion.cpp b/mlir/lib/Linalg/Transforms/Fusion.cpp index ad5bf55..084642e 100644 --- a/mlir/lib/Linalg/Transforms/Fusion.cpp +++ b/mlir/lib/Linalg/Transforms/Fusion.cpp @@ -77,7 +77,7 @@ static llvm::cl::list clTileSizes( // This is achieved by applying the `loopToOperandRangesMaps` permutation maps // to the `loopRanges` in order to obtain view ranges. static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op, - ArrayRef loopRanges, + ArrayRef loopRanges, OperationFolder &state) { ScopedContext scope(b, loc); @@ -92,18 +92,17 @@ static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op, auto map = maps[idx]; LLVM_DEBUG(dbgs() << "map: " << map << "\n"); Value *view = en.value(); - SmallVector viewRanges(map.getNumResults(), nullptr); + SmallVector viewRanges(map.getNumResults()); for (auto en2 : llvm::enumerate(map.getResults())) { unsigned d = en2.index(); // loopToOperandRangesMaps are permutations-only. unsigned loopPos = en2.value().cast().getPosition(); viewRanges[d] = loopRanges[loopPos]; LLVM_DEBUG(dbgs() << "i,j: " << en.index() << ", " << en2.index() << "\t" - << "loopPos: " << loopPos << "\t" << *viewRanges[d]); + << "loopPos: " << loopPos << "\t" << viewRanges[d]); } - // TODO(ntv) opportunities for folding / CSE here rather than build new - // IR (i.e. if dim folds away and the ranges are the same). - clonedViews.push_back(slice(view, viewRanges)); + // TODO(ntv) opportunities for folding/CSE here rather than build new IR. + clonedViews.push_back(b.create(loc, view, viewRanges)); } return op.create(b, loc, clonedViews); } @@ -147,14 +146,13 @@ static Optional fuse(Value *producedView, LinalgOp producer, return llvm::None; unsigned producerIdx = maybeProducerIdx.getValue(); - auto sliceOp = dyn_cast_or_null( + auto sliceOp = dyn_cast_or_null( tiledConsumer.getInput(consumerIdx)->getDefiningOp()); // If we don't have a slice, this also means we don't have loops and the // producer cannot be fused at this level. if (!sliceOp) return llvm::None; - auto sliceRanges = sliceOp.getRanges(); AffineMap producerMap = loopToOperandRangesMaps(producer)[producer.getNumInputs() + producerIdx]; LLVM_DEBUG(dbgs() << "Consumer Idx: " << consumerIdx << "\tmap: " @@ -162,33 +160,30 @@ static Optional fuse(Value *producedView, LinalgOp producer, LLVM_DEBUG(dbgs() << "Producer Idx: " << producerIdx << "\tmap: " << producerMap << "\n"); - SmallVector loopRanges(producer.getNumParallelLoops() + - producer.getNumReductionLoops() + - producer.getNumWindowLoops(), - nullptr); + unsigned nPar = producer.getNumParallelLoops(); + unsigned nRed = producer.getNumReductionLoops(); + unsigned nWin = producer.getNumWindowLoops(); + SmallVector loopRanges(nPar + nRed + nWin); + DenseSet fromSlice; for (auto en : llvm::enumerate(producerMap.getResults())) { - auto *r = *(sliceRanges.begin() + en.index()); // loopToOperandRangesMaps are permutations-only. unsigned posInProducerLoop = en.value().cast().getPosition(); - loopRanges[posInProducerLoop] = r; + loopRanges[posInProducerLoop] = sliceOp.getRange(en.index()); + fromSlice.insert(posInProducerLoop); } OpBuilder b(tiledConsumer.getOperation()); auto loc = tiledConsumer.getLoc(); for (unsigned i = 0; i < loopRanges.size(); ++i) { - auto *r = loopRanges[i]; - if (r) - LLVM_DEBUG(llvm::dbgs() << "LR: " << *r << "\n"); + if (fromSlice.count(i)) + LLVM_DEBUG(llvm::dbgs() << "LR: " << loopRanges[i] << "\n"); else { - // TODO(ntv) opportunities for folding / CSE here rather than build new - // IR (i.e. if dim folds away and the ranges are the same). auto viewDim = getViewDefiningLoopRange(producer, i); - loopRanges[i] = - range(state.create(b, loc, 0), - linalg::intrinsics::dim(viewDim.view, viewDim.dimension), - state.create(b, loc, 1)); - ; - LLVM_DEBUG(llvm::dbgs() << "new LR: " << *loopRanges[i] << "\n"); + loopRanges[i] = SubViewOp::Range{ + state.create(b, loc, 0), + linalg::intrinsics::dim(viewDim.view, viewDim.dimension), + state.create(b, loc, 1)}; + LLVM_DEBUG(llvm::dbgs() << "new LR: " << loopRanges[i] << "\n"); } } diff --git a/mlir/lib/Linalg/Transforms/Tiling.cpp b/mlir/lib/Linalg/Transforms/Tiling.cpp index 0530a81..2b0d739 100644 --- a/mlir/lib/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Linalg/Transforms/Tiling.cpp @@ -116,23 +116,6 @@ static bool isTiledView(LinalgOp linalgOp, unsigned viewIndex, return false; } -static Value *foldRange(Value *view, unsigned dim) { - assert(view->getType().isa() && "View expected"); - if (auto *op = view->getDefiningOp()) { - if (auto viewOp = dyn_cast(op)) { - return *(viewOp.getIndexings().begin() + dim); - } - auto sliceOp = cast(op); - for (auto *i : sliceOp.getIndexings()) - if (i->getType().isa()) { - if (dim == 0) - return i; - --dim; - } - } - return nullptr; -} - static SmallVector makeTiledViews(OpBuilder &b, Location loc, LinalgOp linalgOp, ArrayRef ivs, @@ -162,20 +145,18 @@ static SmallVector makeTiledViews(OpBuilder &b, Location loc, continue; } - // If not a scalar, then construct a new slice. - SmallVector newRanges; - newRanges.reserve(viewRank); + // If not a scalar, then construct a new subview for the tile. + SmallVector subViewOperands; + subViewOperands.reserve(viewRank * 3); for (unsigned r = 0; r < viewRank; ++r) { // Loop position for the range dimension. auto pos = getPosInDomain(linalgOp, viewIndex, r); auto tileSize = tileSizes[pos]; if (isZero(tileSize)) { - auto *foldedRange = foldRange(view, r); - foldedRange ? newRanges.push_back(foldedRange) - : newRanges.push_back( - range(state.create(b, loc, 0), - linalg::intrinsics::dim(view, r), - state.create(b, loc, 1))); + subViewOperands.push_back( + SubViewOp::Range{state.create(b, loc, 0), + linalg::intrinsics::dim(view, r), + state.create(b, loc, 1)}); continue; } @@ -192,16 +173,13 @@ static SmallVector makeTiledViews(OpBuilder &b, Location loc, // Tie this loose end in the future. ValueHandle lb(iv); ValueHandle step(tileSize); - ValueHandle steppedlb = lb + step; - ValueHandle viewSize = linalg::intrinsics::dim(view, r); - ValueHandle ub = select(viewSize < steppedlb, viewSize, steppedlb); + ValueHandle steppedLb = lb + step; // Tiling creates a new slice at the proper index, the slice step is 1 // (i.e. the slice view does not subsample, stepping occurs in the loop). - newRanges.push_back( - range(lb, ub, state.create(b, loc, 1))); + subViewOperands.push_back(SubViewOp::Range{ + iv, steppedLb, state.create(b, loc, 1)}); } - // res.push_back(createOrReturnView(b, loc, viewDefiningOp, newRanges)); - res.push_back(b.create(loc, view, newRanges)); + res.push_back(b.create(loc, view, subViewOperands)); } return res; } diff --git a/mlir/test/Linalg/roundtrip.mlir b/mlir/test/Linalg/roundtrip.mlir index d356c4e..234f82d 100644 --- a/mlir/test/Linalg/roundtrip.mlir +++ b/mlir/test/Linalg/roundtrip.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt %s | mlir-opt | FileCheck %s +// RUN: mlir-opt %s | mlir-opt %s | FileCheck %s // CHECK: #[[map0:.*]] = (d0, d1, d2) -> (d0, d2, d1) // CHECK: #[[map1:.*]] = (d0, d1, d2) -> (d2, d1, d0) diff --git a/mlir/test/Linalg/tile.mlir b/mlir/test/Linalg/tile.mlir index 620ee08..27f5b70 100644 --- a/mlir/test/Linalg/tile.mlir +++ b/mlir/test/Linalg/tile.mlir @@ -10,229 +10,103 @@ // TILE-234-DAG: #[[UB1:.*]] = (d0) -> (d0 + 3) // TILE-234-DAG: #[[UB2:.*]] = (d0) -> (d0 + 4) -func @matmul(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { - %c0 = constant 0 : index - %c1 = constant 1 : index - %I = linalg.range %c0:%arg1:%c1 : !linalg.range - %J = linalg.range %c0:%arg2:%c1 : !linalg.range - %K = linalg.range %c0:%arg3:%c1 : !linalg.range - %A = linalg.view %arg0[%I, %K] : !linalg.view - %B = linalg.view %arg0[%K, %J] : !linalg.view - %C = linalg.view %arg0[%I, %J] : !linalg.view - linalg.matmul(%A, %B, %C) : !linalg.view, !linalg.view, !linalg.view +func @matmul(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { + linalg.matmul(%arg0, %arg1, %arg2) : !linalg.view, !linalg.view, !linalg.view return } -// TILE-2-LABEL: func @matmul(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-2: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-2-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-2-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-2: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view +// TILE-2-LABEL: func @matmul(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-2: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view // TILE-2: linalg.for %i0 = %c0{{.*}} to %[[M]] step %c2 { // TILE-2-NEXT: %[[a:.*]] = affine.apply #[[UB0]](%i0) -// TILE-2-NEXT: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view -// TILE-2-NEXT: %[[cmpuba:.*]] = cmpi "slt", %[[M]], %[[a]] : index -// TILE-2-NEXT: %[[uba:.*]] = select %[[cmpuba]], %[[M]], %[[a]] : index -// TILE-2-NEXT: %[[ra:.*]] = linalg.range %i0:%[[uba]]:%c1 : !linalg.range -// TILE-2: %[[sAi:.*]] = linalg.slice %[[A]][%[[ra]], {{.*}}] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-2: %[[c:.*]] = affine.apply #[[UB0]](%i0) -// TILE-2-NEXT: %[[M:.*]] = linalg.dim %[[C]], 0 : !linalg.view -// TILE-2-NEXT: %[[cmpubc:.*]] = cmpi "slt", %[[M]], %[[c]] : index -// TILE-2-NEXT: %[[ubc:.*]] = select %[[cmpubc]], %[[M]], %[[c]] : index -// TILE-2-NEXT: %[[rc:.*]] = linalg.range %i0:%[[ubc]]:%c1 : !linalg.range -// TILE-2: %[[sCi:.*]] = linalg.slice %[[C]][%[[rc]], {{.*}}] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-2-NEXT: linalg.matmul(%[[sAi]], %[[B]], %[[sCi]]) : !linalg.view, !linalg.view, !linalg.view +// TILE-2-NEXT: %[[K:.*]] = linalg.dim %arg0, 1 : !linalg.view +// TILE-2-NEXT: %[[sAi:.*]] = linalg.subview %arg0[%i0, %[[a]], %c1, %c0, %[[K]], %c1] : !linalg.view +// TILE-2-NEXT: %[[c:.*]] = affine.apply #[[UB0]](%i0) +// TILE-2-NEXT: %[[N:.*]] = linalg.dim %arg2, 1 : !linalg.view +// TILE-2-NEXT: %[[sCi:.*]] = linalg.subview %arg2[%i0, %[[c]], %c1, %c0, %[[N]], %c1] : !linalg.view +// TILE-2-NEXT: linalg.matmul(%[[sAi]], %arg1, %[[sCi]]) : !linalg.view, !linalg.view, !linalg.view -// TILE-02-LABEL: func @matmul(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-02: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-02-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-02-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-02: %[[N:.*]] = linalg.dim %[[B]], 1 : !linalg.view +// TILE-02-LABEL: func @matmul(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-02: %[[N:.*]] = linalg.dim %arg1, 1 : !linalg.view // TILE-02: linalg.for %i0 = %c0 to %[[N]] step %c2 { -// TILE-02: %[[b:.*]] = affine.apply #[[UB0]](%i0) -// TILE-02-NEXT: %[[N:.*]] = linalg.dim %[[B]], 1 : !linalg.view -// TILE-02-NEXT: %[[cmpubb:.*]] = cmpi "slt", %[[N]], %[[b]] : index -// TILE-02-NEXT: %[[ubb:.*]] = select %[[cmpubb]], %[[N]], %[[b]] : index -// TILE-02-NEXT: %[[rb:.*]] = linalg.range %i0:%[[ubb]]:%c1 : !linalg.range -// TILE-02: %[[sBj:.*]] = linalg.slice %[[B]][%{{.*}}, %[[rb]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-02: %[[c:.*]] = affine.apply #[[UB0]](%i0) -// TILE-02: %[[N:.*]] = linalg.dim %[[C]], 1 : !linalg.view -// TILE-02-NEXT: %[[cmpubc:.*]] = cmpi "slt", %[[N]], %[[c]] : index -// TILE-02-NEXT: %[[ubc:.*]] = select %[[cmpubc]], %[[N]], %[[c]] : index -// TILE-02-NEXT: %[[rc:.*]] = linalg.range %i0:%[[ubc]]:%c1 : !linalg.range -// TILE-02-NEXT: %[[sCj:.*]] = linalg.slice %[[C]][%{{.*}}, %[[rc]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-02: linalg.matmul(%[[A]], %[[sBj]], %[[sCj]]) : !linalg.view, !linalg.view, !linalg.view +// TILE-02-NEXT: %[[K:.*]] = linalg.dim %arg1, 0 : !linalg.view +// TILE-02-NEXT: %[[b:.*]] = affine.apply #[[UB0]](%i0) +// TILE-02-NEXT: %[[sBj:.*]] = linalg.subview %arg1[%c0, %[[K]], %c1, %i0, %[[b]], %c1] : !linalg.view +// TILE-02-NEXT: %[[M:.*]] = linalg.dim %arg2, 0 : !linalg.view +// TILE-02-NEXT: %[[c:.*]] = affine.apply #[[UB0]](%i0) +// TILE-02-NEXT: %[[sCj:.*]] = linalg.subview %arg2[%c0, %[[M]], %c1, %i0, %[[c]], %c1] : !linalg.view +// TILE-02-NEXT: linalg.matmul(%arg0, %[[sBj]], %[[sCj]]) : !linalg.view, !linalg.view, !linalg.view -// TILE-002-LABEL: func @matmul(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-002: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-002-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-002-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-002: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view +// TILE-002-LABEL: func @matmul(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-002: %[[K:.*]] = linalg.dim %arg0, 1 : !linalg.view // TILE-002: linalg.for %i0 = %c0{{.*}} to %[[K]] step %c2 { -// TILE-002: %[[a:.*]] = affine.apply #[[UB0]](%i0) -// TILE-002-NEXT: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view -// TILE-002-NEXT: %[[cmpuba:.*]] = cmpi "slt", %[[K]], %[[a]] : index -// TILE-002-NEXT: %[[uba:.*]] = select %[[cmpuba]], %[[K]], %[[a]] : index -// TILE-002-NEXT: %[[ra:.*]] = linalg.range %i0:%[[uba]]:%c1 : !linalg.range -// TILE-002-NEXT: %[[sAj:.*]] = linalg.slice %[[A]][%{{.*}}, %[[ra]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-002: %[[b:.*]] = affine.apply #[[UB0]](%i0) -// TILE-002-NEXT: %[[K:.*]] = linalg.dim %[[B]], 0 : !linalg.view -// TILE-002-NEXT: %[[cmpubb:.*]] = cmpi "slt", %[[K]], %[[b]] : index -// TILE-002-NEXT: %[[ubb:.*]] = select %[[cmpubb]], %[[K]], %[[b]] : index -// TILE-002-NEXT: %[[rb:.*]] = linalg.range %i0:%[[ubb]]:%c1 : !linalg.range -// TILE-002: %[[sBj:.*]] = linalg.slice %[[B]][%[[rb]], %{{.*}}] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-002: linalg.matmul(%[[sAj]], %[[sBj]], %[[C]]) : !linalg.view, !linalg.view, !linalg.view +// TILE-002-NEXT: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view +// TILE-002-NEXT: %[[a:.*]] = affine.apply #[[UB0]](%i0) +// TILE-002-NEXT: %[[sAj:.*]] = linalg.subview %arg0[%c0, %[[M]], %c1, %i0, %[[a]], %c1] : !linalg.view +// TILE-002-NEXT: %[[b:.*]] = affine.apply #[[UB0]](%i0) +// TILE-002-NEXT: %[[N:.*]] = linalg.dim %arg1, 1 : !linalg.view +// TILE-002-NEXT: %[[sBj:.*]] = linalg.subview %arg1[%i0, %[[b]], %c1, %c0, %[[N]], %c1] : !linalg.view +// TILE-002-NEXT: linalg.matmul(%[[sAj]], %[[sBj]], %arg2) : !linalg.view, !linalg.view, !linalg.view -// TILE-234-LABEL: func @matmul(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-234: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-234-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-234-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-234: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view -// TILE-234: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view -// TILE-234: %[[N:.*]] = linalg.dim %[[B]], 1 : !linalg.view +// TILE-234-LABEL: func @matmul(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-234: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view +// TILE-234: %[[K:.*]] = linalg.dim %arg0, 1 : !linalg.view +// TILE-234: %[[N:.*]] = linalg.dim %arg1, 1 : !linalg.view // TILE-234: linalg.for %i0 = %c0{{.*}} to %[[M]] step %c2 { // TILE-234-NEXT: linalg.for %i1 = %c0{{.*}} to %[[N]] step %c3 { // TILE-234-NEXT: linalg.for %i2 = %c0{{.*}} to %[[K]] step %c4 { -// TILE-234-NEXT: %[[ai:.*]] = affine.apply #[[UB0]](%i0) -// TILE-234-NEXT: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubai:.*]] = cmpi "slt", %[[M]], %[[ai]] : index -// TILE-234-NEXT: %[[ubai:.*]] = select %[[cmpubai]], %[[M]], %[[ai]] : index -// TILE-234-NEXT: %[[rai:.*]] = linalg.range %i0:%[[ubai]]:%c1 : !linalg.range -// +// TILE-234-NEXT: %[[ai:.*]] = affine.apply #[[UB0]](%i0) // TILE-234-NEXT: %[[ak:.*]] = affine.apply #[[UB2]](%i2) -// TILE-234-NEXT: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view -// TILE-234-NEXT: %[[cmpubak:.*]] = cmpi "slt", %[[K]], %[[ak]] : index -// TILE-234-NEXT: %[[ubak:.*]] = select %[[cmpubak]], %[[K]], %[[ak]] : index -// TILE-234-NEXT: %[[rak:.*]] = linalg.range %i2:%[[ubak]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sAik:.*]] = linalg.slice %[[A]][%[[rai]], %[[rak]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// +// TILE-234-NEXT: %[[sAik:.*]] = linalg.subview %arg0[%i0, %[[ai]], %c1, %i2, %[[ak]], %c1] : !linalg.view // TILE-234-NEXT: %[[bk:.*]] = affine.apply #[[UB2]](%i2) -// TILE-234-NEXT: %[[K:.*]] = linalg.dim %[[B]], 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubbk:.*]] = cmpi "slt", %[[K]], %[[bk]] : index -// TILE-234-NEXT: %[[ubbk:.*]] = select %[[cmpubbk]], %[[K]], %[[bk]] : index -// TILE-234-NEXT: %[[rbk:.*]] = linalg.range %i2:%[[ubbk]]:%c1 : !linalg.range -// // TILE-234-NEXT: %[[bj:.*]] = affine.apply #[[UB1]](%i1) -// TILE-234-NEXT: %[[N:.*]] = linalg.dim %[[B]], 1 : !linalg.view -// TILE-234-NEXT: %[[cmpubbj:.*]] = cmpi "slt", %[[N]], %[[bj]] : index -// TILE-234-NEXT: %[[ubbj:.*]] = select %[[cmpubbj]], %[[N]], %[[bj]] : index -// TILE-234-NEXT: %[[rbj:.*]] = linalg.range %i1:%[[ubbj]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sBkj:.*]] = linalg.slice %[[B]][%[[rbk]], %[[rbj]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// +// TILE-234-NEXT: %[[sBkj:.*]] = linalg.subview %arg1[%i2, %[[bk]], %c1, %i1, %[[bj]], %c1] : !linalg.view // TILE-234-NEXT: %[[ci:.*]] = affine.apply #[[UB0]](%i0) -// TILE-234-NEXT: %[[M:.*]] = linalg.dim %[[C]], 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubci:.*]] = cmpi "slt", %[[M]], %[[ci]] : index -// TILE-234-NEXT: %[[ubci:.*]] = select %[[cmpubci]], %[[M]], %[[ci]] : index -// TILE-234-NEXT: %[[rci:.*]] = linalg.range %i0:%[[ubci]]:%c1 : !linalg.range -// // TILE-234-NEXT: %[[cj:.*]] = affine.apply #[[UB1]](%i1) -// TILE-234-NEXT: %[[N:.*]] = linalg.dim %[[C]], 1 : !linalg.view -// TILE-234-NEXT: %[[cmpubcj:.*]] = cmpi "slt", %[[N]], %[[cj]] : index -// TILE-234-NEXT: %[[ubcj:.*]] = select %[[cmpubcj]], %[[N]], %[[cj]] : index -// TILE-234-NEXT: %[[rcj:.*]] = linalg.range %i1:%[[ubcj]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sCij:.*]] = linalg.slice %[[C]][%[[rci]], %[[rcj]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view +// TILE-234-NEXT: %[[sCij:.*]] = linalg.subview %arg2[%i0, %[[ci]], %c1, %i1, %[[cj]], %c1] : !linalg.view // // TILE-234-NEXT: linalg.matmul(%[[sAik]], %[[sBkj]], %[[sCij]]) : !linalg.view, !linalg.view, !linalg.view -func @matvec(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { - %c0 = constant 0 : index - %c1 = constant 1 : index - %I = linalg.range %c0:%arg1:%c1 : !linalg.range - %J = linalg.range %c0:%arg2:%c1 : !linalg.range - %2 = linalg.view %arg0[%I, %J] : !linalg.view - %3 = linalg.view %arg0[%J] : !linalg.view - %4 = linalg.view %arg0[%I] : !linalg.view - linalg.matvec(%2, %3, %4) : !linalg.view, !linalg.view, !linalg.view +func @matvec(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { + linalg.matvec(%arg0, %arg1, %arg2) : !linalg.view, !linalg.view, !linalg.view return } -// TILE-2-LABEL: func @matvec(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-2: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-2-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-2-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-2: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view +// TILE-2-LABEL: func @matvec(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-2: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view // TILE-2: linalg.for %i0 = %c0{{.*}} to %[[M]] step %c2 { // TILE-2-NEXT: %[[a:.*]] = affine.apply #[[UB0]](%i0) -// TILE-2-NEXT: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view -// TILE-2-NEXT: %[[cmpuba:.*]] = cmpi "slt", %[[M]], %[[a]] : index -// TILE-2-NEXT: %[[uba:.*]] = select %[[cmpuba]], %[[M]], %[[a]] : index -// TILE-2-NEXT: %[[ra:.*]] = linalg.range %i0:%[[uba]]:%c1 : !linalg.range -// TILE-2: %[[sAi:.*]] = linalg.slice %[[A]][%[[ra]], %{{.*}}] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-2: %[[c:.*]] = affine.apply #[[UB0]](%i0) -// TILE-2-NEXT: %[[M:.*]] = linalg.dim %[[C]], 0 : !linalg.view -// TILE-2-NEXT: %[[cmpubc:.*]] = cmpi "slt", %[[M]], %[[c]] : index -// TILE-2-NEXT: %[[ubc:.*]] = select %[[cmpubc]], %[[M]], %[[c]] : index -// TILE-2-NEXT: %[[rc:.*]] = linalg.range %i0:%[[ubc]]:%c1 : !linalg.range -// TILE-2: %[[sCi:.*]] = linalg.slice %[[C]][%[[rc]]] : !linalg.view, !linalg.range, !linalg.view -// -// TILE-2-NEXT: linalg.matvec(%[[sAi]], %[[B]], %[[sCi]]) : !linalg.view, !linalg.view, !linalg.view +// TILE-2-NEXT: %[[N:.*]] = linalg.dim %arg0, 1 : !linalg.view +// TILE-2-NEXT: %[[sAi:.*]] = linalg.subview %arg0[%i0, %[[a]], %c1, %c0, %[[N]], %c1] : !linalg.view +// TILE-2-NEXT: %[[c:.*]] = affine.apply #[[UB0]](%i0) +// TILE-2-NEXT: %[[sCi:.*]] = linalg.subview %arg2[%i0, %[[c]], %c1] : !linalg.view +// TILE-2-NEXT: linalg.matvec(%[[sAi]], %arg1, %[[sCi]]) : !linalg.view, !linalg.view, !linalg.view -// TILE-02-LABEL: func @matvec(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-02: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-02-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-02-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-02: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view +// TILE-02-LABEL: func @matvec(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-02: %[[K:.*]] = linalg.dim %arg0, 1 : !linalg.view // TILE-02: linalg.for %i0 = %c0{{.*}} to %[[K]] step %c2 { -// TILE-02: %[[a:.*]] = affine.apply #[[UB0]](%i0) -// TILE-02-NEXT: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view -// TILE-02-NEXT: %[[cmpuba:.*]] = cmpi "slt", %[[K]], %[[a]] : index -// TILE-02-NEXT: %[[uba:.*]] = select %[[cmpuba]], %[[K]], %[[a]] : index -// TILE-02-NEXT: %[[ra:.*]] = linalg.range %i0:%[[uba]]:%c1 : !linalg.range -// TILE-02-NEXT: %[[sAj:.*]] = linalg.slice %[[A]][%{{.*}}, %[[ra]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// +// TILE-02-NEXT: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view +// TILE-02-NEXT: %[[a:.*]] = affine.apply #[[UB0]](%i0) +// TILE-02-NEXT: %[[sAj:.*]] = linalg.subview %arg0[%c0, %[[M]], %c1, %i0, %[[a]], %c1] : !linalg.view // TILE-02-NEXT: %[[b:.*]] = affine.apply #[[UB0]](%i0) -// TILE-02-NEXT: %[[K:.*]] = linalg.dim %[[B]], 0 : !linalg.view -// TILE-02-NEXT: %[[cmpubb:.*]] = cmpi "slt", %[[K]], %[[b]] : index -// TILE-02-NEXT: %[[ubb:.*]] = select %[[cmpubb]], %[[K]], %[[b]] : index -// TILE-02-NEXT: %[[rb:.*]] = linalg.range %i0:%[[ubb]]:%c1 : !linalg.range -// TILE-02-NEXT: %[[sBj:.*]] = linalg.slice %[[B]][%[[rb]]] : !linalg.view, !linalg.range, !linalg.view -// -// TILE-02: linalg.matvec(%[[sAj]], %[[sBj]], %[[C]]) : !linalg.view, !linalg.view, !linalg.view +// TILE-02-NEXT: %[[sBj:.*]] = linalg.subview %arg1[%i0, %[[b]], %c1] : !linalg.view +// TILE-02: linalg.matvec(%[[sAj]], %[[sBj]], %arg2) : !linalg.view, !linalg.view, !linalg.view -// TILE-002-LABEL: func @matvec(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { +// TILE-002-LABEL: func @matvec(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { // TILE-002-NOT: linalg.for -// TILE-234-LABEL: func @matvec(%arg0: !linalg.buffer, %arg1: index, %arg2: index, %arg3: index) { -// TILE-234: %[[A:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-234-NEXT: %[[B:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-234-NEXT: %[[C:.*]] = linalg.view %arg0[{{.*}}] : !linalg.view -// TILE-234: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view -// TILE-234: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view +// TILE-234-LABEL: func @matvec(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { +// TILE-234: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view +// TILE-234: %[[K:.*]] = linalg.dim %arg0, 1 : !linalg.view // TILE-234: linalg.for %i0 = %c0{{.*}} to %[[M]] step %c2 { // TILE-234-NEXT: linalg.for %i1 = %c0{{.*}} to %[[K]] step %c3 { // TILE-234-NEXT: %[[ai:.*]] = affine.apply #[[UB0]](%i0) -// TILE-234-NEXT: %[[M:.*]] = linalg.dim %[[A]], 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubai:.*]] = cmpi "slt", %[[M]], %[[ai]] : index -// TILE-234-NEXT: %[[ubai:.*]] = select %[[cmpubai]], %[[M]], %[[ai]] : index -// TILE-234-NEXT: %[[rai:.*]] = linalg.range %i0:%[[ubai]]:%c1 : !linalg.range -// // TILE-234-NEXT: %[[aj:.*]] = affine.apply #[[UB1]](%i1) -// TILE-234-NEXT: %[[K:.*]] = linalg.dim %[[A]], 1 : !linalg.view -// TILE-234-NEXT: %[[cmpubaj:.*]] = cmpi "slt", %[[K]], %[[aj]] : index -// TILE-234-NEXT: %[[ubaj:.*]] = select %[[cmpubaj]], %[[K]], %[[aj]] : index -// TILE-234-NEXT: %[[raj:.*]] = linalg.range %i1:%[[ubaj]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sAij:.*]] = linalg.slice %[[A]][%[[rai]], %[[raj]]] : !linalg.view, !linalg.range, !linalg.range, !linalg.view -// -// TILE-234-NEXT: %[[b:.*]] = affine.apply #[[UB1]](%i1) -// TILE-234-NEXT: %[[K:.*]] = linalg.dim %[[B]], 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubb:.*]] = cmpi "slt", %[[K]], %[[b]] : index -// TILE-234-NEXT: %[[ubb:.*]] = select %[[cmpubb]], %[[K]], %[[b]] : index -// TILE-234-NEXT: %[[rb:.*]] = linalg.range %i1:%[[ubb]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sB:.*]] = linalg.slice %[[B]][%[[rb]]] : !linalg.view, !linalg.range, !linalg.view +// TILE-234-NEXT: %[[sAij:.*]] = linalg.subview %arg0[%i0, %[[ai]], %c1, %i1, %[[aj]], %c1] : !linalg.view +// TILE-234-NEXT: %[[bj:.*]] = affine.apply #[[UB1]](%i1) +// TILE-234-NEXT: %[[sBj:.*]] = linalg.subview %arg1[%i1, %[[bj]], %c1] : !linalg.view +// TILE-234-NEXT: %[[ci:.*]] = affine.apply #[[UB0]](%i0) +// TILE-234-NEXT: %[[sCi:.*]] = linalg.subview %arg2[%i0, %[[ci]], %c1] : !linalg.view // -// TILE-234-NEXT: %[[c:.*]] = affine.apply #[[UB0]](%i0) -// TILE-234-NEXT: %[[M:.*]] = linalg.dim %[[C]], 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubc:.*]] = cmpi "slt", %[[M]], %[[c]] : index -// TILE-234-NEXT: %[[ubc:.*]] = select %[[cmpubc]], %[[M]], %[[c]] : index -// TILE-234-NEXT: %[[rc:.*]] = linalg.range %i0:%[[ubc]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sC:.*]] = linalg.slice %[[C]][%[[rc]]] : !linalg.view, !linalg.range, !linalg.view -// -// TILE-234-NEXT: linalg.matvec(%[[sAij]], %[[sB]], %[[sC]]) : !linalg.view, !linalg.view, !linalg.view +// TILE-234-NEXT: linalg.matvec(%[[sAij]], %[[sBj]], %[[sCi]]) : !linalg.view, !linalg.view, !linalg.view func @dot(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { linalg.dot(%arg0, %arg1, %arg2) : !linalg.view, !linalg.view, !linalg.view @@ -242,19 +116,9 @@ func @dot(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg // TILE-2: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view // TILE-2: linalg.for %i0 = %c0{{.*}} to %[[M]] step %c2 { // TILE-2-NEXT: %[[a:.*]] = affine.apply #[[UB0]](%i0) -// TILE-2-NEXT: %[[M:.*]] = linalg.dim %arg0, 0 : !linalg.view -// TILE-2-NEXT: %[[cmpuba:.*]] = cmpi "slt", %[[M]], %[[a]] : index -// TILE-2-NEXT: %[[uba:.*]] = select %[[cmpuba]], %[[M]], %[[a]] : index -// TILE-2-NEXT: %[[ra:.*]] = linalg.range %i0:%[[uba]]:%c1 : !linalg.range -// TILE-2: %[[sAi:.*]] = linalg.slice %arg0[%[[ra]]] : !linalg.view, !linalg.range, !linalg.view -// +// TILE-2-NEXT: %[[sAi:.*]] = linalg.subview %arg0[%i0, %[[a]], %c1] : !linalg.view // TILE-2-NEXT: %[[b:.*]] = affine.apply #[[UB0]](%i0) -// TILE-2-NEXT: %[[K:.*]] = linalg.dim %arg1, 0 : !linalg.view -// TILE-2-NEXT: %[[cmpubb:.*]] = cmpi "slt", %[[K]], %[[b]] : index -// TILE-2-NEXT: %[[ubb:.*]] = select %[[cmpubb]], %[[K]], %[[b]] : index -// TILE-2-NEXT: %[[rb:.*]] = linalg.range %i0:%[[ubb]]:%c1 : !linalg.range -// TILE-2-NEXT: %[[sBi:.*]] = linalg.slice %arg1[%[[rb]]] : !linalg.view, !linalg.range, !linalg.view -// +// TILE-2-NEXT: %[[sBi:.*]] = linalg.subview %arg1[%i0, %[[b]], %c1] : !linalg.view // TILE-2-NEXT: linalg.dot(%[[sAi]], %[[sBi]], {{.*}}) : !linalg.view, !linalg.view, !linalg.view // TILE-02-LABEL: func @dot(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg.view) { @@ -267,17 +131,7 @@ func @dot(%arg0: !linalg.view, %arg1: !linalg.view, %arg2: !linalg // TILE-234: %[[K:.*]] = linalg.dim %arg0, 0 : !linalg.view // TILE-234: linalg.for %i0 = %c0{{.*}} to %[[K]] step %c2 { // TILE-234-NEXT: %[[a:.*]] = affine.apply #[[UB0]](%i0) -// TILE-234-NEXT: %[[K:.*]] = linalg.dim %arg0, 0 : !linalg.view -// TILE-234-NEXT: %[[cmpuba:.*]] = cmpi "slt", %[[K]], %[[a]] : index -// TILE-234-NEXT: %[[uba:.*]] = select %[[cmpuba]], %[[K]], %[[a]] : index -// TILE-234-NEXT: %[[ra:.*]] = linalg.range %i0:%[[uba]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sA:.*]] = linalg.slice %arg0[%[[ra]]] : !linalg.view, !linalg.range, !linalg.view -// +// TILE-234-NEXT: %[[sAi:.*]] = linalg.subview %arg0[%i0, %[[a]], %c1] : !linalg.view // TILE-234-NEXT: %[[b:.*]] = affine.apply #[[UB0]](%i0) -// TILE-234-NEXT: %[[K:.*]] = linalg.dim %arg1, 0 : !linalg.view -// TILE-234-NEXT: %[[cmpubb:.*]] = cmpi "slt", %[[K]], %[[b]] : index -// TILE-234-NEXT: %[[ubb:.*]] = select %[[cmpubb]], %[[K]], %[[b]] : index -// TILE-234-NEXT: %[[rb:.*]] = linalg.range %i0:%[[ubb]]:%c1 : !linalg.range -// TILE-234-NEXT: %[[sB:.*]] = linalg.slice %arg1[%[[rb]]] : !linalg.view, !linalg.range, !linalg.view -// -// TILE-234-NEXT: linalg.dot(%[[sA]], %[[sB]], %arg2) : !linalg.view, !linalg.view, !linalg.view +// TILE-234-NEXT: %[[sBi:.*]] = linalg.subview %arg1[%i0, %[[b]], %c1] : !linalg.view +// TILE-234-NEXT: linalg.dot(%[[sAi]], %[[sBi]], %arg2) : !linalg.view, !linalg.view, !linalg.view -- 2.7.4