diff --git a/lib/Dialect/TMTensor/Transforms/Bufferize.cpp b/lib/Dialect/TMTensor/Transforms/Bufferize.cpp index f70099e0a478..dd8218b54f57 100644 --- a/lib/Dialect/TMTensor/Transforms/Bufferize.cpp +++ b/lib/Dialect/TMTensor/Transforms/Bufferize.cpp @@ -32,18 +32,8 @@ namespace mlir::torch::TMTensor { #define GEN_PASS_DEF_TMTENSORBUFFERIZE #include "torch-mlir-dialects/Dialect/TMTensor/Transforms/Passes.h.inc" -static Value cloneMemref(Location loc, Value memref, OpBuilder &b) { - auto memrefType = cast(memref.getType()); - auto alloc = - memref::AllocOp::create(b, loc, memref::getMixedSizes(b, loc, memref), - memrefType.getElementType()); - memref::CopyOp::create(b, loc, memref, alloc); - return alloc; -} - static LogicalResult allocateBuffersForResults(Location loc, TMTensorOp tmtensorOp, - ValueRange outputs, SmallVectorImpl &resultBuffers, OpBuilder &b) { // Lazily compute loopRanges. SmallVector loopRanges; @@ -62,12 +52,23 @@ allocateBuffersForResults(Location loc, TMTensorOp tmtensorOp, } auto tensorShape = tensorType.getShape(); auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType()); - Value resultTensor = outputs[resultIndex]; // Clone output buffers whose value is actually used. OpOperand *tiedOpOperand = tmtensorOp.getOutputOperand(resultIndex); if (tmtensorOp.payloadUsesValueFromOperand(tiedOpOperand)) { - resultBuffers.push_back(cloneMemref(loc, resultTensor, b)); + // The op updates this operand in place, so the buffer it writes has to + // start out holding the operand's incoming value. Copy from the *tensor* + // operand, not from the buffer the conversion driver handed us: a copy + // with a tensor source is itself the read of that tensor, so the read + // lands here, at the op, instead of at the operand's definition -- the + // same reason the `ins` crossings below are built here. + Value alloc = memref::AllocOp::create( + b, loc, tensor::getMixedSizes(b, loc, tiedOpOperand->get()), + memrefType.getElementType()); + bufferization::MaterializeInDestinationOp::create( + b, loc, /*result=*/TypeRange{}, tiedOpOperand->get(), alloc, + /*restrict=*/true, /*writable=*/true); + resultBuffers.push_back(alloc); continue; } @@ -77,8 +78,13 @@ allocateBuffersForResults(Location loc, TMTensorOp tmtensorOp, continue; } + // The payload does not read this output, so the buffer needs the operand's + // shape but not its contents. Take the dynamic sizes from the *tensor* + // operand; the destination-passing-style contract guarantees they match the + // result's -- "Init operands and their tied OpResults have the same type. + // Dynamic dimension sizes also match at runtime." resultBuffers.push_back(memref::AllocOp::create( - b, loc, memref::getMixedSizes(b, loc, resultTensor), + b, loc, tensor::getMixedSizes(b, loc, tiedOpOperand->get()), memrefType.getElementType())); } return success(); @@ -108,16 +114,37 @@ class BufferizeAnyTMTensorOp : public OpInterfaceConversionPattern { Location loc = op.getLoc(); SmallVector newOutputBuffers; - SmallVector outputs(operands.begin() + op.getNumInputs(), - operands.end()); - if (failed(allocateBuffersForResults(loc, op, outputs, newOutputBuffers, - rewriter))) { + if (failed( + allocateBuffersForResults(loc, op, newOutputBuffers, rewriter))) { return op.emitOpError() << "Failed to allocate buffers for tensor results."; } - SmallVector inputs(operands.begin(), - operands.begin() + op.getNumInputs()); + // Take each tensor `ins` operand's buffer here, at the op that reads it, + // instead of using the operand the conversion driver converted for us. The + // driver pins a target materialization to the operand's *definition*, which + // leaves the read invisible to a later full bufferization: this pass is + // partial, so the TMTensor op it feeds reads raw memory, and One-Shot + // Bufferize analyses tensor operands only. Any write between the definition + // and this op may then be bufferized in place over the very buffer we were + // handed, silently changing the value the op observes. Stating the crossing + // here instead puts the read at the program point where it happens, so + // One-Shot sees the read-after-write and preserves the value. `read_only` + // is truthful -- a TMTensor op never writes an `ins` operand -- and keeps + // the crossing from counting as a write as well. + SmallVector inputs; + for (OpOperand &opOperand : + op->getOpOperands().take_front(op.getNumInputs())) { + auto tensorType = dyn_cast(opOperand.get().getType()); + if (!tensorType) { + inputs.push_back(operands[opOperand.getOperandNumber()]); + continue; + } + inputs.push_back(bufferization::ToBufferOp::create( + rewriter, loc, + MemRefType::get(tensorType.getShape(), tensorType.getElementType()), + opOperand.get(), /*read_only=*/true)); + } createTMTensorOpOnBuffers(rewriter, op, inputs, newOutputBuffers); // Replace the results of the old op with the new output buffers. rewriter.replaceOp(op, newOutputBuffers); @@ -131,7 +158,8 @@ static Value materializeToTensor(OpBuilder &builder, TensorType type, ValueRange inputs, Location loc) { assert(inputs.size() == 1); assert(isa(inputs[0].getType())); - return bufferization::ToTensorOp::create(builder, loc, type, inputs[0]); + return bufferization::ToTensorOp::create(builder, loc, type, inputs[0], + /*restrict=*/true); } /// Converts TMTensor operations that work on tensor-type operands or results to @@ -188,8 +216,9 @@ struct TMTensorBufferizePass }); // Mark all Standard operations legal. - target.addLegalDialect(); + target.addLegalDialect< + arith::ArithDialect, bufferization::BufferizationDialect, + func::FuncDialect, memref::MemRefDialect, tensor::TensorDialect>(); // Mark all TMTensor operations illegal as long as they work on tensors. auto isLegalOperation = [&](Operation *op) { diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index dbfa86aa5aa7..d33816d1b25a 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -16,6 +16,13 @@ set(TORCH_MLIR_TEST_DEPENDS TorchMLIRPythonModules torch-mlir-opt torch-mlir-capi-torch-test + # Needed by tests whose RUN lines reach for upstream tooling: + # --one-shot-bufferize is not registered in torch-mlir-opt, and the + # numeric regression tests execute the lowered IR. + mlir-opt + mlir-runner + mlir_runner_utils + mlir_c_runner_utils ) add_lit_testsuite(check-torch-mlir "Running the torch-mlir regression tests" diff --git a/test/Dialect/TMTensor/bufferize-one-shot-hazard.mlir b/test/Dialect/TMTensor/bufferize-one-shot-hazard.mlir new file mode 100644 index 000000000000..1074406935b9 --- /dev/null +++ b/test/Dialect/TMTensor/bufferize-one-shot-hazard.mlir @@ -0,0 +1,89 @@ +// RUN: torch-mlir-opt %s -tm-tensor-bufferize -tm-tensor-to-loops \ +// RUN: | mlir-opt --one-shot-bufferize="allow-unknown-ops" \ +// RUN: | FileCheck %s + +// The IR shape `convert-torch-to-tmtensor` emits for +// `torch.aten.scatter_reduce.two` with `include_self=false`: one `tensor.empty` +// init with two writers -- the reduction identity and the update values -- read +// by two separate `tm_tensor.scatter`s, a "clear" pass and an "add" pass. +// +// The property pinned here: the two loops must read from two *different* +// buffers, and the buffer the clear loop reads must be the one filled with 0.0. +// See bufferize-scatter-reduce-numeric.mlir for the same bug at the level of +// values. + +// CHECK-LABEL: func.func @scatter_reduce_sum_no_include_self( +// CHECK-SAME: %[[INDICES:.*]]: tensor<2x1xi32>, +// CHECK-SAME: %[[SRC:.*]]: tensor<2xf32>, +// CHECK-SAME: %[[SELF:.*]]: tensor<4xf32> +// CHECK-DAG: %[[SELF_BUF:.*]] = bufferization.to_buffer %[[SELF]] +// CHECK-DAG: %[[SRC_BUF:.*]] = bufferization.to_buffer %[[SRC]] +// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 + +// The two writers of the shared init get separate buffers: neither can be +// written in place, because the scatter that reads the other one conflicts. +// CHECK: %[[VALS:.*]] = memref.alloc(){{.*}} : memref<2xf32> +// CHECK: %[[ZEROS:.*]] = memref.alloc(){{.*}} : memref<2xf32> +// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[ZEROS]] : memref<2xf32>) +// CHECK: linalg.generic {{.*}} ins(%[[SRC_BUF]]{{.*}}) outs(%[[VALS]] : memref<2xf32>) + +// The clear loop reads the zeros and overwrites its targets. +// CHECK: %[[CLEARED:.*]] = memref.alloc() : memref<4xf32> +// CHECK: memref.copy %[[SELF_BUF]], %[[CLEARED]] +// CHECK: %[[IDX_BUF:.*]] = bufferization.to_buffer %[[INDICES]] read_only +// CHECK: scf.for +// CHECK: %[[U0:.*]] = memref.load %[[ZEROS]] +// CHECK: memref.load %[[IDX_BUF]] +// CHECK: memref.store %[[U0]], %[[CLEARED]] +// CHECK: } + +// The add loop reads the values -- a *different* buffer. +// CHECK: %[[RESULT:.*]] = memref.alloc() : memref<4xf32> +// CHECK: %[[RESULT_TENSOR:.*]] = bufferization.to_tensor %[[RESULT]] restrict +// CHECK: memref.copy %[[CLEARED]], %[[RESULT]] +// CHECK: %[[IDX_BUF_2:.*]] = bufferization.to_buffer %[[INDICES]] read_only +// CHECK: scf.for +// CHECK: %[[U1:.*]] = memref.load %[[VALS]] +// CHECK: memref.load %[[IDX_BUF_2]] +// CHECK: %[[O:.*]] = memref.load %[[RESULT]] +// CHECK: %[[SUM:.*]] = arith.addf %[[U1]], %[[O]] : f32 +// CHECK: memref.store %[[SUM]], %[[RESULT]] +// CHECK: } +// CHECK: return %[[RESULT_TENSOR]] : tensor<4xf32> +func.func @scatter_reduce_sum_no_include_self( + %indices: tensor<2x1xi32>, %src: tensor<2xf32>, + %self: tensor<4xf32>) -> tensor<4xf32> { + %zero = arith.constant 0.000000e+00 : f32 + %init = tensor.empty() : tensor<2xf32> + + // Writer 1: the reduction identity. + %zeros = linalg.fill ins(%zero : f32) outs(%init : tensor<2xf32>) -> tensor<2xf32> + + // Writer 2: the update values, sharing the SAME init tensor. + %vals = linalg.generic { + indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], + iterator_types = ["parallel"]} + ins(%src : tensor<2xf32>) outs(%init : tensor<2xf32>) { + ^bb0(%in: f32, %out: f32): + linalg.yield %in : f32 + } -> tensor<2xf32> + + // Clear pass: plant the identity at every scatter target. + %cleared = tm_tensor.scatter {dimension_map = array} unique_indices(true) + ins(%zeros, %indices : tensor<2xf32>, tensor<2x1xi32>) + outs(%self : tensor<4xf32>) { + ^bb0(%update: f32, %orig: f32): + tm_tensor.yield %update : f32 + } -> tensor<4xf32> + + // Add pass: accumulate the values on top. + %result = tm_tensor.scatter {dimension_map = array} unique_indices(false) + ins(%vals, %indices : tensor<2xf32>, tensor<2x1xi32>) + outs(%cleared : tensor<4xf32>) { + ^bb0(%update: f32, %orig: f32): + %sum = arith.addf %update, %orig : f32 + tm_tensor.yield %sum : f32 + } -> tensor<4xf32> + + return %result : tensor<4xf32> +} diff --git a/test/Dialect/TMTensor/bufferize-scatter-reduce-numeric.mlir b/test/Dialect/TMTensor/bufferize-scatter-reduce-numeric.mlir new file mode 100644 index 000000000000..fae068b718fe --- /dev/null +++ b/test/Dialect/TMTensor/bufferize-scatter-reduce-numeric.mlir @@ -0,0 +1,145 @@ +// REQUIRES: mlir-runner +// RUN: torch-mlir-opt %s \ +// RUN: --pass-pipeline='builtin.module(func.func(convert-torch-to-tmtensor,convert-torch-to-linalg,canonicalize,cse,torch-finalizing-backend-type-conversion,tm-tensor-bufferize,tm-tensor-to-loops))' \ +// RUN: | mlir-opt --one-shot-bufferize="bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map" \ +// RUN: --convert-linalg-to-loops --expand-strided-metadata --lower-affine \ +// RUN: --convert-scf-to-cf --finalize-memref-to-llvm --convert-func-to-llvm \ +// RUN: --convert-arith-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts \ +// RUN: | mlir-runner -e main --entry-point-result=void \ +// RUN: --shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \ +// RUN: --shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext \ +// RUN: | FileCheck %s + +// Numeric regression test for `torch.aten.scatter_reduce.two` with +// `include_self=false`, starting from the actual torch op and *executing* the +// result. This is the only level at which the bug this file guards against is +// observable as a value rather than as a shape of IR. +// +// The pipeline deliberately finishes bufferization with plain +// `--one-shot-bufferize`, so that One-Shot's analysis actually runs. Adding +// `copy-before-write` -- as `refbackend`'s pipeline does -- would skip the +// analysis and copy defensively, which hides the bug. +// +// An e2e test cannot stand in for this file. `ScatterReduceFloatSumModule` +// already *is* this op with `include_self=false`, and it passes on a build +// without the fix, for two independent reasons. First, `refbackend`'s pipeline +// passes `copy-before-write`, which skips One-Shot's analysis altogether. +// Second, it would pass even with the analysis enabled, because the pipeline +// opens with `linalg-fuse-elementwise-ops`, whose `RemoveOutsDependency` +// pattern rewrites every `outs` operand the payload does not read to a fresh +// `tensor.empty` -- dissolving the shared init described below before +// bufferization is ever consulted. +// +// Neither defence is one this pass can rely on. `copy-before-write` is opt-in, +// and the `RemoveOutsDependency` rewrite is undone by any `cse` between fusion +// and bufferization: the fresh `tensor.empty` is structurally identical to the +// one the `linalg.fill` below already writes into, so CSE merges the two and +// the shared init comes straight back. What keeps the e2e suite green is +// therefore an accident of pass ordering, not anything the suite can be +// configured to expose -- which is why this file pins the behaviour at the +// level of values instead. +// +// The hazard `tm-tensor-bufferize` is handed: `convert-torch-to-tmtensor` +// computes the reduction identity (a `linalg.fill`) and the update values (a +// `linalg.generic`) into the *same* `tensor.empty` init, then runs two +// `tm_tensor.scatter`s -- a "clear" pass that plants the identity at each target +// and a "reduce" pass that combines the updates in. One init, several writers. +// +// `cse` in the pipeline above is load-bearing, not tidying: the lowering emits +// two structurally identical `tensor.empty` + `linalg.fill` pairs, and it is CSE +// that merges them into the single init the two scatters then share. Drop it and +// this test passes even on a build without the fix. +// +// If the scatters' reads are stated at their operands' definitions rather than +// at the scatters, One-Shot Bufferize sees no conflict between those writers and +// folds them onto one buffer, so one of the two values is lost. See +// bufferize-one-shot-hazard.mlir for the IR-level property. +// +// Both modes below use the same inputs: +// +// self = [10 20; 30 40; 50 60; 70 80] +// index = [[0 1]; [0 3]] src = [[1 2]; [3 4]] +// scatter_reduce(dim=0, include_self=false) +// +// Cell (0,0) receives both 1 and 3; (1,1) receives 2; (3,1) receives 4; every +// other cell keeps its `self` value. +// +// Two reduce modes are checked because the collapse destroys a different value +// in each, and a fix that recovered only one shape would still be wrong: +// +// sum -- the identity is 0.0, which is also the constant the `linalg.generic` +// init is filled with, so CSE merges the two fills into one that +// *precedes* the generic. The surviving value is the updates, which +// the clear pass then scatters instead of zeros and the add pass adds +// a second time: 2x at every singly-targeted cell. (At (0,0) it is +// last-write plus full sum, 3 + (1+3) = 7, not 8.) +// +// amax -- the identity is -inf, a different constant, so its fill is a third +// writer of the shared init emitted *after* the generic. That fill is +// what survives, destroying the updates, and both passes read -inf. +// +// mode correct collapsed +// sum [4 20; 30 2; 50 60; 70 4] [7 20; 30 4; 50 60; 70 8] +// amax [3 20; 30 2; 50 60; 70 4] [-inf 20; 30 -inf; 50 60; 70 -inf] +// +// CHECK: Unranked Memref +// CHECK-NEXT: [4, 20] +// CHECK-NEXT: [30, 2] +// CHECK-NEXT: [50, 60] +// CHECK-NEXT: [70, 4] +// CHECK: Unranked Memref +// CHECK-NEXT: [3, 20] +// CHECK-NEXT: [30, 2] +// CHECK-NEXT: [50, 60] +// CHECK-NEXT: [70, 4] + +func.func private @printMemrefF32(tensor<*xf32>) + +// The `torch_c` casts wrap the torch op inside a builtin-tensor signature so +// that @main can call it directly and the whole test stays in one file. They +// are erased by `torch-finalizing-backend-type-conversion`. +func.func @scatter_reduce_sum(%self_b: tensor<4x2xf32>, %index_b: tensor<2x2xi64>, + %src_b: tensor<2x2xf32>) -> tensor<4x2xf32> { + %self = torch_c.from_builtin_tensor %self_b : tensor<4x2xf32> -> !torch.vtensor<[4,2],f32> + %index = torch_c.from_builtin_tensor %index_b : tensor<2x2xi64> -> !torch.vtensor<[2,2],si64> + %src = torch_c.from_builtin_tensor %src_b : tensor<2x2xf32> -> !torch.vtensor<[2,2],f32> + %dim = torch.constant.int 0 + %reduce = torch.constant.str "sum" + %include_self = torch.constant.bool false + %0 = torch.aten.scatter_reduce.two %self, %dim, %index, %src, %reduce, %include_self : + !torch.vtensor<[4,2],f32>, !torch.int, !torch.vtensor<[2,2],si64>, + !torch.vtensor<[2,2],f32>, !torch.str, !torch.bool -> !torch.vtensor<[4,2],f32> + %r = torch_c.to_builtin_tensor %0 : !torch.vtensor<[4,2],f32> -> tensor<4x2xf32> + return %r : tensor<4x2xf32> +} + +func.func @scatter_reduce_amax(%self_b: tensor<4x2xf32>, %index_b: tensor<2x2xi64>, + %src_b: tensor<2x2xf32>) -> tensor<4x2xf32> { + %self = torch_c.from_builtin_tensor %self_b : tensor<4x2xf32> -> !torch.vtensor<[4,2],f32> + %index = torch_c.from_builtin_tensor %index_b : tensor<2x2xi64> -> !torch.vtensor<[2,2],si64> + %src = torch_c.from_builtin_tensor %src_b : tensor<2x2xf32> -> !torch.vtensor<[2,2],f32> + %dim = torch.constant.int 0 + %reduce = torch.constant.str "amax" + %include_self = torch.constant.bool false + %0 = torch.aten.scatter_reduce.two %self, %dim, %index, %src, %reduce, %include_self : + !torch.vtensor<[4,2],f32>, !torch.int, !torch.vtensor<[2,2],si64>, + !torch.vtensor<[2,2],f32>, !torch.str, !torch.bool -> !torch.vtensor<[4,2],f32> + %r = torch_c.to_builtin_tensor %0 : !torch.vtensor<[4,2],f32> -> tensor<4x2xf32> + return %r : tensor<4x2xf32> +} + +func.func @main() { + %self = arith.constant dense<[[10.0, 20.0], [30.0, 40.0], [50.0, 60.0], [70.0, 80.0]]> : tensor<4x2xf32> + %index = arith.constant dense<[[0, 1], [0, 3]]> : tensor<2x2xi64> + %src = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32> + %sum = call @scatter_reduce_sum(%self, %index, %src) + : (tensor<4x2xf32>, tensor<2x2xi64>, tensor<2x2xf32>) -> tensor<4x2xf32> + %sum_unranked = tensor.cast %sum : tensor<4x2xf32> to tensor<*xf32> + call @printMemrefF32(%sum_unranked) : (tensor<*xf32>) -> () + + %amax = call @scatter_reduce_amax(%self, %index, %src) + : (tensor<4x2xf32>, tensor<2x2xi64>, tensor<2x2xf32>) -> tensor<4x2xf32> + %amax_unranked = tensor.cast %amax : tensor<4x2xf32> to tensor<*xf32> + call @printMemrefF32(%amax_unranked) : (tensor<*xf32>) -> () + return +} diff --git a/test/Dialect/TMTensor/bufferize.mlir b/test/Dialect/TMTensor/bufferize.mlir index 0f04dc5197e0..c543630517e6 100644 --- a/test/Dialect/TMTensor/bufferize.mlir +++ b/test/Dialect/TMTensor/bufferize.mlir @@ -1,16 +1,25 @@ // RUN: torch-mlir-opt -split-input-file -tm-tensor-bufferize %s | FileCheck %s +// Each tensor operand crosses to a buffer at the TMTensor op that reads it, not +// at the operand's definition, so a later full bufferization can see the read. +// See bufferize-one-shot-hazard.mlir for the end-to-end consequence. + // ----- // CHECK-LABEL: func.func @scan_1d_inclusive( // CHECK-SAME: %[[IN_TENSOR:.*]]: tensor<128xi32>, %[[OUT_TENSOR:.*]]: tensor<128xi32>, // CHECK-SAME: %[[ACC_TENSOR:.*]]: tensor) -> (tensor<128xi32>, tensor) { -// CHECK-DAG: %[[IN_MEMREF:.*]] = bufferization.to_buffer %[[IN_TENSOR]] : tensor<128xi32> to memref<128xi32> -// CHECK-DAG: %[[OUT_MEMREF_NEW:.*]] = memref.alloc() : memref<128xi32> -// CHECK-DAG: %[[ACC_MEMREF_NEW:.*]] = memref.alloc() : memref -// CHECK-DAG: %[[OUT_TENSOR_NEW:.*]] = bufferization.to_tensor %[[OUT_MEMREF_NEW]] : memref<128xi32> to tensor<128xi32> -// CHECK-DAG: %[[ACC_TENSOR_NEW:.*]] = bufferization.to_tensor %[[ACC_MEMREF_NEW]] : memref to tensor -// CHECK: tm_tensor.scan dimension(0) inclusive(true) ins(%[[IN_MEMREF]] : memref<128xi32>) -// CHECK-SAME: outs(%[[OUT_MEMREF_NEW]], %[[ACC_MEMREF_NEW]] : memref<128xi32>, memref) { +// The payload reads neither `out` nor `acc`, so neither incoming value is +// copied. CHECK-NEXT from here to the op is what makes that absence +// enforceable: a bare CHECK-NOT would only constrain the gap between the two +// positive matches bracketing it, leaving the rest of the prologue free to hold +// a copy. +// CHECK-NEXT: %[[OUT_MEMREF:.*]] = memref.alloc() : memref<128xi32> +// CHECK-NEXT: %[[OUT_TENSOR_NEW:.*]] = bufferization.to_tensor %[[OUT_MEMREF]] restrict : memref<128xi32> to tensor<128xi32> +// CHECK-NEXT: %[[ACC_MEMREF:.*]] = memref.alloc() : memref +// CHECK-NEXT: %[[ACC_TENSOR_NEW:.*]] = bufferization.to_tensor %[[ACC_MEMREF]] restrict : memref to tensor +// CHECK-NEXT: %[[IN_MEMREF:.*]] = bufferization.to_buffer %[[IN_TENSOR]] read_only : tensor<128xi32> to memref<128xi32> +// CHECK-NEXT: tm_tensor.scan dimension(0) inclusive(true) ins(%[[IN_MEMREF]] : memref<128xi32>) +// CHECK-SAME: outs(%[[OUT_MEMREF]], %[[ACC_MEMREF]] : memref<128xi32>, memref) { // CHECK: ^bb0(%[[OUT_PREV_ELEMENT:.*]]: i32, %[[IN_ELEMENT:.*]]: i32): // CHECK: %[[OUT_CURRENT_ELEMENT:.*]] = arith.addi %[[OUT_PREV_ELEMENT]], %[[IN_ELEMENT]] : i32 // CHECK: tm_tensor.yield %[[OUT_CURRENT_ELEMENT]] : i32 @@ -30,15 +39,16 @@ func.func @scan_1d_inclusive(%in: tensor<128xi32>, %out: tensor<128xi32>, %acc: // CHECK-LABEL: func.func @scan_1d_exclusive( // CHECK-SAME: %[[IN_TENSOR:.*]]: tensor<128xi32>, %[[OUT_TENSOR:.*]]: tensor<128xi32>, // CHECK-SAME: %[[ACC_TENSOR:.*]]: tensor) -> (tensor<128xi32>, tensor) { -// CHECK-DAG: %[[IN_MEMREF:.*]] = bufferization.to_buffer %[[IN_TENSOR]] : tensor<128xi32> to memref<128xi32> -// CHECK-DAG: %[[ACC_MEMREF:.*]] = bufferization.to_buffer %[[ACC_TENSOR]] : tensor to memref -// CHECK-DAG: %[[OUT_MEMREF_NEW:.*]] = memref.alloc() : memref<128xi32> -// CHECK-DAG: %[[ACC_MEMREF_NEW:.*]] = memref.alloc() : memref -// CHECK-DAG: %[[OUT_TENSOR_NEW:.*]] = bufferization.to_tensor %[[OUT_MEMREF_NEW]] : memref<128xi32> to tensor<128xi32> -// CHECK-DAG: %[[ACC_TENSOR_NEW:.*]] = bufferization.to_tensor %[[ACC_MEMREF_NEW]] : memref to tensor -// CHECK: memref.copy %[[ACC_MEMREF]], %[[ACC_MEMREF_NEW]] : memref to memref +// CHECK: %[[OUT_MEMREF:.*]] = memref.alloc() : memref<128xi32> +// CHECK: %[[OUT_TENSOR_NEW:.*]] = bufferization.to_tensor %[[OUT_MEMREF]] restrict : memref<128xi32> to tensor<128xi32> +// CHECK: %[[ACC_MEMREF:.*]] = memref.alloc() : memref +// CHECK: %[[ACC_TENSOR_NEW:.*]] = bufferization.to_tensor %[[ACC_MEMREF]] restrict : memref to tensor +// The payload reads `acc`, so its incoming value is copied into the buffer the +// op updates in place -- straight from the tensor operand, so exactly once. +// CHECK: bufferization.materialize_in_destination %[[ACC_TENSOR]] in restrict writable %[[ACC_MEMREF]] : (tensor, memref) -> () +// CHECK: %[[IN_MEMREF:.*]] = bufferization.to_buffer %[[IN_TENSOR]] read_only : tensor<128xi32> to memref<128xi32> // CHECK: tm_tensor.scan dimension(0) inclusive(false) ins(%[[IN_MEMREF]] : memref<128xi32>) -// CHECK-SAME: outs(%[[OUT_MEMREF_NEW]], %[[ACC_MEMREF_NEW]] : memref<128xi32>, memref) { +// CHECK-SAME: outs(%[[OUT_MEMREF]], %[[ACC_MEMREF]] : memref<128xi32>, memref) { // CHECK: ^bb0(%[[OUT_PREV_ELEMENT:.*]]: i32, %[[IN_ELEMENT:.*]]: i32): // CHECK: %[[OUT_CURRENT_ELEMENT:.*]] = arith.addi %[[OUT_PREV_ELEMENT]], %[[IN_ELEMENT]] : i32 // CHECK: tm_tensor.yield %[[OUT_CURRENT_ELEMENT]] : i32 @@ -54,19 +64,46 @@ func.func @scan_1d_exclusive(%in: tensor<128xi32>, %out: tensor<128xi32>, %acc: return %ret_out, %ret_acc: tensor<128xi32>, tensor } +// ----- +// CHECK-LABEL: func.func @scan_1d_dynamic( +// CHECK-SAME: %[[IN_TENSOR:.*]]: tensor, %[[OUT_TENSOR:.*]]: tensor, +// CHECK-SAME: %[[ACC_TENSOR:.*]]: tensor) -> (tensor, tensor) { +// The size comes off the tensor argument, NOT off a buffer, and neither `outs` +// operand is read by the payload, so nothing is copied anywhere. CHECK-NEXT +// from here to the op is what makes those absences enforceable -- it leaves no +// gap for a `memref.dim` on a buffer, or for a copy, to hide in. +// CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index +// CHECK-NEXT: %[[OUT_DIM:.*]] = tensor.dim %[[OUT_TENSOR]], %[[C0]] : tensor +// CHECK-NEXT: %[[OUT_MEMREF:.*]] = memref.alloc(%[[OUT_DIM]]) : memref +// CHECK-NEXT: %[[OUT_TENSOR_NEW:.*]] = bufferization.to_tensor %[[OUT_MEMREF]] restrict : memref to tensor +// CHECK-NEXT: %[[ACC_MEMREF:.*]] = memref.alloc() : memref +// CHECK-NEXT: %[[ACC_TENSOR_NEW:.*]] = bufferization.to_tensor %[[ACC_MEMREF]] restrict : memref to tensor +// CHECK-NEXT: %[[IN_MEMREF:.*]] = bufferization.to_buffer %[[IN_TENSOR]] read_only : tensor to memref +// CHECK-NEXT: tm_tensor.scan dimension(0) inclusive(true) ins(%[[IN_MEMREF]] : memref) +// CHECK-SAME: outs(%[[OUT_MEMREF]], %[[ACC_MEMREF]] : memref, memref) { +// CHECK: return %[[OUT_TENSOR_NEW]], %[[ACC_TENSOR_NEW]] : tensor, tensor +func.func @scan_1d_dynamic(%in: tensor, %out: tensor, %acc: tensor) -> (tensor, tensor) { + %ret_out, %ret_acc = tm_tensor.scan dimension(0) inclusive(true) + ins(%in : tensor) outs(%out, %acc: tensor, tensor) { + ^bb0(%arg0 : i32, %arg1 : i32): + %sum = arith.addi %arg0, %arg1 : i32 + tm_tensor.yield %sum : i32 + } -> tensor, tensor + return %ret_out, %ret_acc: tensor, tensor +} + // ----- // CHECK-LABEL: func.func @scatter_update_scalar_1D( // CHECK-SAME: %[[ORIG_TENSOR:.*]]: tensor<8xi32>, // CHECK-SAME: %[[INDICES_TENSOR:.*]]: tensor<3x1xi32>, // CHECK-SAME: %[[UPDATES_TENSOR:.*]]: tensor<3xi32>) -> tensor<8xi32> { -// CHECK-DAG: %[[UPDATES_MEMREF:.*]] = bufferization.to_buffer %[[UPDATES_TENSOR]] : tensor<3xi32> to memref<3xi32> -// CHECK-DAG: %[[INDICES_MEMREF:.*]] = bufferization.to_buffer %[[INDICES_TENSOR]] : tensor<3x1xi32> to memref<3x1xi32> -// CHECK-DAG: %[[ORIG_MEMREF:.*]] = bufferization.to_buffer %[[ORIG_TENSOR]] : tensor<8xi32> to memref<8xi32> -// CHECK-DAG: %[[ORIG_MEMREF_NEW:.*]] = memref.alloc() : memref<8xi32> -// CHECK-DAG: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[ORIG_MEMREF_NEW]] : memref<8xi32> to tensor<8xi32> -// CHECK: memref.copy %[[ORIG_MEMREF]], %[[ORIG_MEMREF_NEW]] : memref<8xi32> to memref<8xi32> +// CHECK: %[[ORIG_MEMREF:.*]] = memref.alloc() : memref<8xi32> +// CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[ORIG_MEMREF]] restrict : memref<8xi32> to tensor<8xi32> +// CHECK: bufferization.materialize_in_destination %[[ORIG_TENSOR]] in restrict writable %[[ORIG_MEMREF]] : (tensor<8xi32>, memref<8xi32>) -> () +// CHECK: %[[UPDATES_MEMREF:.*]] = bufferization.to_buffer %[[UPDATES_TENSOR]] read_only : tensor<3xi32> to memref<3xi32> +// CHECK: %[[INDICES_MEMREF:.*]] = bufferization.to_buffer %[[INDICES_TENSOR]] read_only : tensor<3x1xi32> to memref<3x1xi32> // CHECK: tm_tensor.scatter {dimension_map = array} unique_indices(true) ins(%[[UPDATES_MEMREF]], %[[INDICES_MEMREF]] -// CHECK-SAME: : memref<3xi32>, memref<3x1xi32>) outs(%[[ORIG_MEMREF_NEW]] : memref<8xi32>) { +// CHECK-SAME: : memref<3xi32>, memref<3x1xi32>) outs(%[[ORIG_MEMREF]] : memref<8xi32>) { // CHECK: ^bb0(%[[UPDATE_SCALAR:.*]]: i32, %[[ORIG_SCALAR:.*]]: i32): // CHECK: tm_tensor.yield %[[UPDATE_SCALAR]] : i32 // CHECK: } @@ -87,14 +124,13 @@ func.func @scatter_update_scalar_1D( // CHECK-SAME: %[[ORIG_TENSOR:.*]]: tensor<8xi32>, // CHECK-SAME: %[[INDICES_TENSOR:.*]]: tensor<3x1xi32>, // CHECK-SAME: %[[UPDATES_TENSOR:.*]]: tensor<3xi32>) -> tensor<8xi32> { -// CHECK-DAG: %[[UPDATES_MEMREF:.*]] = bufferization.to_buffer %[[UPDATES_TENSOR]] : tensor<3xi32> to memref<3xi32> -// CHECK-DAG: %[[INDICES_MEMREF:.*]] = bufferization.to_buffer %[[INDICES_TENSOR]] : tensor<3x1xi32> to memref<3x1xi32> -// CHECK-DAG: %[[ORIG_MEMREF:.*]] = bufferization.to_buffer %[[ORIG_TENSOR]] : tensor<8xi32> to memref<8xi32> -// CHECK-DAG: %[[ORIG_MEMREF_NEW:.*]] = memref.alloc() : memref<8xi32> -// CHECK-DAG: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[ORIG_MEMREF_NEW]] : memref<8xi32> to tensor<8xi32> -// CHECK: memref.copy %[[ORIG_MEMREF]], %[[ORIG_MEMREF_NEW]] : memref<8xi32> to memref<8xi32> +// CHECK: %[[ORIG_MEMREF:.*]] = memref.alloc() : memref<8xi32> +// CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[ORIG_MEMREF]] restrict : memref<8xi32> to tensor<8xi32> +// CHECK: bufferization.materialize_in_destination %[[ORIG_TENSOR]] in restrict writable %[[ORIG_MEMREF]] : (tensor<8xi32>, memref<8xi32>) -> () +// CHECK: %[[UPDATES_MEMREF:.*]] = bufferization.to_buffer %[[UPDATES_TENSOR]] read_only : tensor<3xi32> to memref<3xi32> +// CHECK: %[[INDICES_MEMREF:.*]] = bufferization.to_buffer %[[INDICES_TENSOR]] read_only : tensor<3x1xi32> to memref<3x1xi32> // CHECK: tm_tensor.scatter {dimension_map = array} unique_indices(true) ins(%[[UPDATES_MEMREF]], %[[INDICES_MEMREF]] -// CHECK-SAME: : memref<3xi32>, memref<3x1xi32>) outs(%[[ORIG_MEMREF_NEW]] : memref<8xi32>) { +// CHECK-SAME: : memref<3xi32>, memref<3x1xi32>) outs(%[[ORIG_MEMREF]] : memref<8xi32>) { // CHECK: ^bb0(%[[UPDATE_SCALAR:.*]]: i32, %[[ORIG_SCALAR:.*]]: i32): // CHECK: %[[CST1:.*]] = arith.constant 1 : i32 // CHECK: %[[ADD:.*]] = arith.addi %[[ORIG_SCALAR]], %[[CST1]] : i32 diff --git a/test/lit.cfg.py b/test/lit.cfg.py index dcc730360983..6323b9b65ed4 100644 --- a/test/lit.cfg.py +++ b/test/lit.cfg.py @@ -34,12 +34,21 @@ config.substitutions.append(("%PATH%", config.environment["PATH"])) config.substitutions.append(("%shlibext", config.llvm_shlib_ext)) +config.substitutions.append(("%mlir_lib_dir", config.llvm_shlib_dir)) # Register optional-backend availability as lit features so tests that exercise a # specific backend contract can guard themselves with `// REQUIRES: `. if getattr(config, "enable_stablehlo", False): config.available_features.add("stablehlo") +# Tests that execute IR need `mlir-runner` plus the runner support libraries. A +# build configured for torch-mlir alone need not have produced them, so gate +# those tests with `// REQUIRES: mlir-runner` rather than failing. +if os.path.isfile( + os.path.join(config.llvm_shlib_dir, "libmlir_runner_utils" + config.llvm_shlib_ext) +) and lit.util.which("mlir-runner", config.llvm_tools_dir): + config.available_features.add("mlir-runner") + llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"]) # llvm_config.use_default_substitutions() @@ -97,6 +106,12 @@ ] tools = [ "torch-mlir-opt", + # Used by tests that need an upstream pass torch-mlir-opt does not register, + # e.g. --one-shot-bufferize. + "mlir-opt", + # Used by tests that execute the lowered IR. Guarded by `REQUIRES: + # mlir-runner`. + "mlir-runner", ToolSubst("%PYTHON", config.python_executable, unresolved="ignore"), ]