From 0f8da4a5c977b646ad2f16a0774efc1e44b42e2a Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Fri, 7 Feb 2025 17:43:11 -0800 Subject: [PATCH] update for comments Signed-off-by: Rob Suderman --- .../compiler/Dialect/Stream/IR/StreamOps.cpp | 33 ------------------- .../compiler/Dialect/Stream/IR/StreamOps.td | 9 ++--- .../Stream/Transforms/ExecutionPlacement.cpp | 11 ++++--- .../Transforms/test/schedule_execution.mlir | 8 ++--- 4 files changed, 13 insertions(+), 48 deletions(-) diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp index 4ee287e3deb52..775a256d4d050 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp +++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp @@ -2495,39 +2495,6 @@ LogicalResult AsyncTransferOp::verify() { return success(); } -IREE::Stream::AffinityAttr AsyncTransferOp::getAffinityAttr() { - if (getExecAffinityAttr()) { - return getExecAffinityAttr(); - } - - auto sourceType = cast(getSource().getType()); - auto resultType = cast(getResult().getType()); - if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging && - resultType.getLifetime() == IREE::Stream::Lifetime::Staging) { - // TODO(multi-device): figure out how to model staging->staging transfers. - return getSourceAffinityAttr(); - } else if (sourceType.getLifetime() == IREE::Stream::Lifetime::External || - sourceType.getLifetime() == IREE::Stream::Lifetime::Staging) { - // If source is staging then the op should execute on the consumer. - return getResultAffinityAttr(); - } else if (resultType.getLifetime() == IREE::Stream::Lifetime::External || - resultType.getLifetime() == IREE::Stream::Lifetime::Staging) { - // If result is staging then the op should execute on the producer. - return getSourceAffinityAttr(); - } else { - // Default to result affinity. - return getSourceAffinityAttr(); - } -} - -void AsyncTransferOp::setAffinityAttr(IREE::Stream::AffinityAttr value) { - if (value) { - setExecAffinityAttr(value); - } else { - removeExecAffinityAttr(); - } -} - void AsyncTransferOp::build(OpBuilder &builder, OperationState &state, Type type, Value source, Value source_size, Value result_size, AffinityAttr source_attr, diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td index 276b54b0cea96..1999283d0711f 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td +++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td @@ -2344,10 +2344,7 @@ def Stream_AsyncBarrierOp : Stream_Op<"async.barrier", [ } def Stream_AsyncTransferOp : Stream_Op<"async.transfer", [ - DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, Stream_AsyncPhaseOp, Stream_StreamableOp, DeclareOpInterfaceMethods:$source_affinity, OptionalAttr:$result_affinity, - OptionalAttr:$exec_affinity + OptionalAttr:$affinity ); let results = (outs AnyTypeOf<[ @@ -2384,10 +2381,10 @@ def Stream_AsyncTransferOp : Stream_Op<"async.transfer", [ let assemblyFormat = [{ $source `:` type($source) `` `{` $source_size `}` + (`on` `(` $affinity^ `)`)? (`from` `(` $source_affinity^ `)`)? `->` (`to` `(` $result_affinity^ `)`)? - (`on` `(` $exec_affinity^ `)`)? type($result) `` `{` $result_size `}` attr-dict-with-keyword }]; diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ExecutionPlacement.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ExecutionPlacement.cpp index b459595aa633c..3b399b2920bcc 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ExecutionPlacement.cpp +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ExecutionPlacement.cpp @@ -43,21 +43,22 @@ struct ExecutionPlacementPass void runOnOperation() override { getOperation()->walk([](IREE::Stream::AsyncTransferOp transfer) { - if (transfer.getExecAffinityAttr()) + if (transfer.getAffinityAttr()) return; auto operand = transfer.getSource(); auto producer = operand.getDefiningOp(); auto streamable = dyn_cast_or_null(producer); - auto srcAffinity = dyn_cast(producer); + auto srcAffinity = + dyn_cast_or_null(producer); bool hasOneUse = operand.hasOneUse(); if (hasOneUse && streamable && srcAffinity) { - transfer.setExecAffinityAttr(srcAffinity.getAffinityAttr()); - } else { - transfer.setExecAffinityAttr(transfer.getResultAffinityAttr()); + transfer.setAffinityAttr(srcAffinity.getAffinityAttr()); + return; } + transfer.setAffinityAttr(transfer.getResultAffinityAttr()); }); } }; diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir index b1dfdca62561d..f9d68b194997f 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir @@ -100,14 +100,14 @@ util.func public @deviceTripleSync(%arg0: i1) -> (!stream.resource, ! // CHECK: stream.async.dispatch // CHECK: stream.async.transfer %3 = stream.async.dispatch on(#hal.device.affinity<@device1>) @ex::@dispatch0[%c1, %c1, %c1](%0[%c0 to %c128 for %c128]) : (!stream.resource{%c128}) -> !stream.resource{%c128} - %4 = stream.async.transfer %3 : !stream.resource{%c128} from(#hal.device.affinity<@device1>) -> to(#hal.device.affinity<@device0>) on(#hal.device.affinity<@device1>) !stream.resource{%c128} + %4 = stream.async.transfer %3 : !stream.resource{%c128} on(#hal.device.affinity<@device1>) from(#hal.device.affinity<@device1>) -> to(#hal.device.affinity<@device0>) !stream.resource{%c128} // CHECK: stream.async.execute // CHECK: stream.async.splat // CHECK: stream.async.dispatch // CHECK: stream.async.transfer %5 = stream.async.dispatch on(#hal.device.affinity<@device2>) @ex::@dispatch0[%c1, %c1, %c1](%0[%c0 to %c128 for %c128]) : (!stream.resource{%c128}) -> !stream.resource{%c128} - %6 = stream.async.transfer %5 : !stream.resource{%c128} from(#hal.device.affinity<@device2>) -> to(#hal.device.affinity<@device0>) on(#hal.device.affinity<@device2>) !stream.resource{%c128} + %6 = stream.async.transfer %5 : !stream.resource{%c128} on(#hal.device.affinity<@device2>) from(#hal.device.affinity<@device2>) -> to(#hal.device.affinity<@device0>) !stream.resource{%c128} // CHECK: stream.async.execute // CHECK: stream.async.dispatch @@ -121,13 +121,13 @@ util.func public @deviceTripleSync(%arg0: i1) -> (!stream.resource, ! // CHECK: stream.async.execute // CHECK: stream.async.transfer // CHECK: stream.async.dispatch - %9 = stream.async.transfer %7 : !stream.resource{%c128} from(#hal.device.affinity<@device0>) -> to(#hal.device.affinity<@device1>) on(#hal.device.affinity<@device1>) !stream.resource{%c128} + %9 = stream.async.transfer %7 : !stream.resource{%c128} on(#hal.device.affinity<@device1>) from(#hal.device.affinity<@device0>) -> to(#hal.device.affinity<@device1>) !stream.resource{%c128} %12 = stream.async.dispatch on(#hal.device.affinity<@device1>) @ex::@dispatch0[%c1, %c1, %c1](%9[%c0 to %c128 for %c128]) : (!stream.resource{%c128}) -> !stream.resource{%c128} // CHECK: stream.async.execute // CHECK: stream.async.transfer // CHECK: stream.async.dispatch - %10 = stream.async.transfer %7 : !stream.resource{%c128} from(#hal.device.affinity<@device0>) -> to(#hal.device.affinity<@device2>) on(#hal.device.affinity<@device2>) !stream.resource{%c128} + %10 = stream.async.transfer %7 : !stream.resource{%c128} on(#hal.device.affinity<@device2>) from(#hal.device.affinity<@device0>) -> to(#hal.device.affinity<@device2>) !stream.resource{%c128} %13 = stream.async.dispatch on(#hal.device.affinity<@device2>) @ex::@dispatch0[%c1, %c1, %c1](%10[%c0 to %c128 for %c128]) : (!stream.resource{%c128}) -> !stream.resource{%c128} util.return %11, %12, %13 : !stream.resource, !stream.resource, !stream.resource