diff --git a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt index 5b0cb6c37bfab5..4f8382e8e6e6b3 100644 --- a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt @@ -1,3 +1,5 @@ +add_subdirectory(TransformOps) + add_mlir_dialect(DLTI dlti) add_mlir_doc(DLTIAttrs DLTIDialect Dialects/ -gen-dialect-doc) diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h index f50a654f3885d5..6e2623e084be52 100644 --- a/mlir/include/mlir/Dialect/DLTI/DLTI.h +++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h @@ -22,6 +22,19 @@ namespace detail { class DataLayoutEntryAttrStorage; } // namespace detail } // namespace mlir +namespace mlir { +namespace dlti { +/// Find the first DataLayoutSpec associated to `op`, via either the +/// DataLayoutOpInterface, a method on ModuleOp, or an attribute implementing +/// the interface, on `op` and else on `op`'s ancestors in turn. +DataLayoutSpecInterface getDataLayoutSpec(Operation *op); + +/// Find the first TargetSystemSpec associated to `op`, via either the +/// DataLayoutOpInterface, a method on ModuleOp, or an attribute implementing +/// the interface, on `op` and else on `op`'s ancestors in turn. +TargetSystemSpecInterface getTargetSystemSpec(Operation *op); +} // namespace dlti +} // namespace mlir #define GET_ATTRDEF_CLASSES #include "mlir/Dialect/DLTI/DLTIAttrs.h.inc" diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt new file mode 100644 index 00000000000000..1188d1a9b24d60 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt @@ -0,0 +1,6 @@ +set(LLVM_TARGET_DEFINITIONS DLTITransformOps.td) +mlir_tablegen(DLTITransformOps.h.inc -gen-op-decls) +mlir_tablegen(DLTITransformOps.cpp.inc -gen-op-defs) +add_public_tablegen_target(MLIRDLTITransformOpsIncGen) + +add_mlir_doc(DLTITransformOps DLTITransformOps Dialects/ -gen-op-doc) diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h new file mode 100644 index 00000000000000..970f2597ef7ad3 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h @@ -0,0 +1,38 @@ +//===- DLTITransformOps.h - DLTI transform ops ------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_DLTI_TRANSFORMOPS_DLTITRANSFORMOPS_H +#define MLIR_DIALECT_DLTI_TRANSFORMOPS_DLTITRANSFORMOPS_H + +#include "mlir/Dialect/Transform/IR/TransformAttrs.h" +#include "mlir/Dialect/Transform/IR/TransformDialect.h" +#include "mlir/Dialect/Transform/IR/TransformTypes.h" +#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h" + +namespace mlir { +namespace transform { +class QueryOp; +} // namespace transform +} // namespace mlir + +namespace mlir { +class DialectRegistry; + +namespace dlti { +void registerTransformDialectExtension(DialectRegistry ®istry); +} // namespace dlti +} // namespace mlir + +////===----------------------------------------------------------------------===// +//// DLTI Transform Operations +////===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h.inc" + +#endif // MLIR_DIALECT_DLTI_TRANSFORMOPS_DLTITRANSFORMOPS_H diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td new file mode 100644 index 00000000000000..69aacac986ad73 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td @@ -0,0 +1,61 @@ +//===- DLTITransformOps.td - DLTI transform ops ------------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef DLTI_TRANSFORM_OPS +#define DLTI_TRANSFORM_OPS + +include "mlir/Dialect/Transform/IR/TransformDialect.td" +include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td" +include "mlir/Dialect/Transform/IR/TransformTypes.td" +include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/IR/OpBase.td" + +def QueryOp : Op +]> { + let summary = "Return attribute (as param) associated to key via DTLI"; + let description = [{ + This op queries data layout and target information associated to payload + IR by way of the DLTI dialect. A lookup is performed for the given `key` + at the `target` op, with the DLTI dialect determining which interfaces and + attributes are consulted - first checking `target` and then its ancestors. + + When only `key` is provided, the lookup occurs with respect to the data + layout specification of DLTI. When `device` is provided, the lookup occurs + with respect to DLTI's target device specifications associated to a DLTI + system device specification. + + #### Return modes + + When succesful, the result, `associated_attr`, associates one attribute as a + param for each op in `target`'s payload. + + If the lookup fails - as DLTI specifications or entries with the right + names are missing (i.e. the values of `device` and `key`) - a definite + failure is returned. + }]; + + let arguments = (ins TransformHandleTypeInterface:$target, + OptionalAttr:$device, + StrAttr:$key); + let results = (outs TransformParamTypeInterface:$associated_attr); + let assemblyFormat = + "(`:``:` $device^ `:``:`)? $key `at` $target attr-dict `:`" + "functional-type(operands, results)"; + + let extraClassDeclaration = [{ + ::mlir::DiagnosedSilenceableFailure applyToOne( + ::mlir::transform::TransformRewriter &rewriter, + ::mlir::Operation *target, + ::mlir::transform::ApplyToEachResultList &results, + TransformState &state); + }]; +} + +#endif // DLTI_TRANSFORM_OPS diff --git a/mlir/include/mlir/InitAllExtensions.h b/mlir/include/mlir/InitAllExtensions.h index 20a4ab6f18a286..0adc5e52f2a0e5 100644 --- a/mlir/include/mlir/InitAllExtensions.h +++ b/mlir/include/mlir/InitAllExtensions.h @@ -25,6 +25,7 @@ #include "mlir/Conversion/UBToLLVM/UBToLLVM.h" #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h" #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h" +#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h" #include "mlir/Dialect/Func/Extensions/AllExtensions.h" #include "mlir/Dialect/Func/TransformOps/FuncTransformOps.h" #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h" @@ -69,6 +70,7 @@ inline void registerAllExtensions(DialectRegistry ®istry) { // Register all transform dialect extensions. affine::registerTransformDialectExtension(registry); bufferization::registerTransformDialectExtension(registry); + dlti::registerTransformDialectExtension(registry); func::registerTransformDialectExtension(registry); gpu::registerTransformDialectExtension(registry); linalg::registerTransformDialectExtension(registry); diff --git a/mlir/lib/Dialect/DLTI/CMakeLists.txt b/mlir/lib/Dialect/DLTI/CMakeLists.txt index f74f5f80130908..7691a8a10a37ad 100644 --- a/mlir/lib/Dialect/DLTI/CMakeLists.txt +++ b/mlir/lib/Dialect/DLTI/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(TransformOps) add_mlir_dialect_library(MLIRDLTIDialect DLTI.cpp Traits.cpp diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp index 420c605d1a19b2..c995f44c380e57 100644 --- a/mlir/lib/Dialect/DLTI/DLTI.cpp +++ b/mlir/lib/Dialect/DLTI/DLTI.cpp @@ -393,6 +393,41 @@ TargetSystemSpecAttr::verify(function_ref emitError, // DLTIDialect //===----------------------------------------------------------------------===// +DataLayoutSpecInterface dlti::getDataLayoutSpec(Operation *op) { + DataLayoutSpecInterface dlSpec = nullptr; + + for (Operation *cur = op; cur && !dlSpec; cur = cur->getParentOp()) { + if (auto dataLayoutOp = dyn_cast(cur)) + dlSpec = dataLayoutOp.getDataLayoutSpec(); + else if (auto moduleOp = dyn_cast(cur)) + dlSpec = moduleOp.getDataLayoutSpec(); + else + for (NamedAttribute attr : cur->getAttrs()) + if ((dlSpec = llvm::dyn_cast(attr.getValue()))) + break; + } + + return dlSpec; +} + +TargetSystemSpecInterface dlti::getTargetSystemSpec(Operation *op) { + TargetSystemSpecInterface sysSpec = nullptr; + + for (Operation *cur = op; cur && !sysSpec; cur = cur->getParentOp()) { + if (auto dataLayoutOp = dyn_cast(cur)) + sysSpec = dataLayoutOp.getTargetSystemSpec(); + else if (auto moduleOp = dyn_cast(cur)) + sysSpec = moduleOp.getTargetSystemSpec(); + else + for (NamedAttribute attr : cur->getAttrs()) + if ((sysSpec = + llvm::dyn_cast(attr.getValue()))) + break; + } + + return sysSpec; +} + constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName; constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessKey; constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessBig; diff --git a/mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt new file mode 100644 index 00000000000000..05734353f9d11d --- /dev/null +++ b/mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt @@ -0,0 +1,15 @@ +add_mlir_dialect_library(MLIRDLTITransformOps + DLTITransformOps.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/DLTI/TransformOps + + DEPENDS + MLIRDLTITransformOpsIncGen + MLIRDLTIDialect + + LINK_LIBS PUBLIC + MLIRDLTIDialect + MLIRSideEffectInterfaces + MLIRTransformDialect + ) diff --git a/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp new file mode 100644 index 00000000000000..259af8857c02f6 --- /dev/null +++ b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp @@ -0,0 +1,94 @@ + +//===- DLTITransformOps.cpp - Implementation of DLTI transform ops --------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h" + +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h" +#include "mlir/Dialect/Transform/Utils/Utils.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" + +using namespace mlir; +using namespace mlir::transform; + +#define DEBUG_TYPE "dlti-transforms" + +//===----------------------------------------------------------------------===// +// QueryOp +//===----------------------------------------------------------------------===// + +void transform::QueryOp::getEffects( + SmallVectorImpl &effects) { + onlyReadsHandle(getTargetMutable(), effects); + producesHandle(getOperation()->getOpResults(), effects); + onlyReadsPayload(effects); +} + +DiagnosedSilenceableFailure transform::QueryOp::applyToOne( + transform::TransformRewriter &rewriter, Operation *target, + transform::ApplyToEachResultList &results, TransformState &state) { + StringAttr deviceId = getDeviceAttr(); + StringAttr key = getKeyAttr(); + + DataLayoutEntryInterface entry; + if (deviceId) { + TargetSystemSpecInterface sysSpec = dlti::getTargetSystemSpec(target); + if (!sysSpec) + return mlir::emitDefiniteFailure(target->getLoc()) + << "no target system spec associated to: " << target; + + if (auto targetSpec = sysSpec.getDeviceSpecForDeviceID(deviceId)) + entry = targetSpec->getSpecForIdentifier(key); + else + return mlir::emitDefiniteFailure(target->getLoc()) + << "no " << deviceId << " target device spec found"; + } else { + DataLayoutSpecInterface dlSpec = dlti::getDataLayoutSpec(target); + if (!dlSpec) + return mlir::emitDefiniteFailure(target->getLoc()) + << "no data layout spec associated to: " << target; + + entry = dlSpec.getSpecForIdentifier(key); + } + + if (!entry) + return mlir::emitDefiniteFailure(target->getLoc()) + << "no DLTI entry for key: " << key; + + results.push_back(entry.getValue()); + + return DiagnosedSilenceableFailure::success(); +} + +//===----------------------------------------------------------------------===// +// Transform op registration +//===----------------------------------------------------------------------===// + +namespace { +class DLTITransformDialectExtension + : public transform::TransformDialectExtension< + DLTITransformDialectExtension> { +public: + using Base::Base; + + void init() { + registerTransformOps< +#define GET_OP_LIST +#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.cpp.inc" + >(); + } +}; +} // namespace + +#define GET_OP_CLASSES +#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.cpp.inc" + +void mlir::dlti::registerTransformDialectExtension(DialectRegistry ®istry) { + registry.addExtensions(); +} diff --git a/mlir/test/Dialect/DLTI/query.mlir b/mlir/test/Dialect/DLTI/query.mlir new file mode 100644 index 00000000000000..2504958de09289 --- /dev/null +++ b/mlir/test/Dialect/DLTI/query.mlir @@ -0,0 +1,296 @@ +// RUN: mlir-opt -transform-interpreter -canonicalize -split-input-file -verify-diagnostics %s | FileCheck %s + +// expected-remark @below {{associated attr 42 : i32}} +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %module = transform.get_parent_op %funcs : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query "test.id" at %module : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %module : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + // expected-remark @below {{associated attr 24 : i32}} + func.func private @f() attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 24 : i32>>} +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query "test.id" at %funcs : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %funcs : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +// expected-remark @below {{associated attr 42 : i32}} +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + func.func private @f() attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 24 : i32>>} +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %module = transform.get_parent_op %funcs : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query "test.id" at %module : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %module : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + func.func @matmul_tensors( + %arg0: tensor, %arg1: tensor, %arg2: tensor) + -> tensor { + // expected-remark @below {{associated attr 42 : i32}} + %0 = linalg.matmul ins(%arg0, %arg1: tensor, tensor) + outs(%arg2: tensor) + -> tensor + return %0 : tensor + } +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query "test.id" at %matmul : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %matmul : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + func.func @matmul_tensors( + %arg0: tensor, %arg1: tensor, %arg2: tensor) + -> tensor attributes {test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 24 : i32>>} { + // expected-remark @below {{associated attr 24 : i32}} + %0 = linalg.matmul ins(%arg0, %arg1: tensor, tensor) + outs(%arg2: tensor) + -> tensor + return %0 : tensor + } +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query "test.id" at %matmul : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %matmul : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +// expected-remark @below {{associated attr 42 : i32}} +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>>} { + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %module = transform.get_parent_op %func : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"CPU"::"test.id" at %module : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %module : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>, + "GPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 43 : i32>>>} { + // expected-remark @below {{associated attr 43 : i32}} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"GPU"::"test.id" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>, + "GPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 43 : i32>>>} { + // expected-remark @below {{associated attr 24 : i32}} + func.func private @f() attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 24 : i32>>> } +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"CPU"::"test.id" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"cache::L1::size_in_bytes", 65536 : i32>, + #dlti.dl_entry<"cache::L1d::size_in_bytes", 32768 : i32>>> } { + // expected-remark @below {{L1::size_in_bytes 65536 : i32}} + // expected-remark @below {{L1d::size_in_bytes 32768 : i32}} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %l1_size = transform.dlti.query ::"CPU"::"cache::L1::size_in_bytes" at %func : (!transform.any_op) -> !transform.param + %l1d_size = transform.dlti.query ::"CPU"::"cache::L1d::size_in_bytes" at %func : (!transform.any_op) -> !transform.param + transform.debug.emit_param_as_remark %l1_size, "L1::size_in_bytes" at %func : !transform.param, !transform.any_op + transform.debug.emit_param_as_remark %l1d_size, "L1d::size_in_bytes" at %func : !transform.param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": + #dlti.target_device_spec<#dlti.dl_entry<"inner_most_tile_size", 42 : i32>>>} { + // CHECK-LABEL: func @matmul_tensors + func.func @matmul_tensors( + %arg0: tensor, %arg1: tensor, %arg2: tensor) + -> tensor { + // CHECK: scf.for {{.*}} to {{.*}} step {{.*}}42 + // CHECK: tensor.extract_slice + // CHECK: linalg.matmul + // CHECK: tensor.insert_slice + // CHECK: scf.yield + %0 = linalg.matmul ins(%arg0, %arg1: tensor, tensor) + outs(%arg2: tensor) + -> tensor + // CHECK: return + return %0 : tensor + } +} + +// Demonstrates obtaining transform op parameters from DLTI attributes and directly putting them to use. +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg : (!transform.any_op) -> !transform.any_op + %tile_size = transform.dlti.query ::"CPU"::"inner_most_tile_size" at %matmul : (!transform.any_op) -> !transform.param + transform.structured.tile_using_for %matmul tile_sizes [%tile_size] : (!transform.any_op, !transform.param) -> (!transform.any_op, !transform.any_op) + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>, + "GPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 43 : i32>>>} { + // expected-error @below {{no "NPU" target device spec found}} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"NPU"::"test.id" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>, + "GPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 43 : i32>>>} { + // expected-error @below {{no DLTI entry for key: "unspecified"}} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"CPU"::"unspecified" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>, + "GPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 43 : i32>>>} { + // expected-error @below {{no data layout spec associated to: }} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query "test.id" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + // expected-error @below {{no target system spec associated to: }} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"CPU"::"test.id" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module { + // expected-error @below {{no target system spec associated to: }} + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + %param = transform.dlti.query ::"CPU"::"test.id" at %func : (!transform.any_op) -> !transform.any_param + transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op + transform.yield + } +} + +// ----- + +module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} { + func.func private @f() +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg: !transform.any_op) { + %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op + // expected-error @below {{expected the type of the parameter attribute ('i32') to match the parameter type ('i64')}} + %param = transform.dlti.query "test.id" at %funcs : (!transform.any_op) -> !transform.param + transform.debug.emit_param_as_remark %param, "associated attr" at %funcs : !transform.param, !transform.any_op + transform.yield + } +}