Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds UnitAttribute to root ops for tuner #19345

Merged
merged 8 commits into from
Dec 4, 2024
Merged

Adds UnitAttribute to root ops for tuner #19345

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
14e9c96
Add `UnitAttribute` to root ops for tuner
nithinsubbiah Dec 2, 2024
b1e8726
Add lit test
nithinsubbiah Dec 2, 2024
01ff091
Fix whitespace
nithinsubbiah Dec 2, 2024
8654394
Update LLVMCPU lit test
nithinsubbiah Dec 3, 2024
53c15ff
Add test case in strip-compilation-info pass
nithinsubbiah Dec 3, 2024
9d91a9d
Guard tuner attributes with flag
nithinsubbiah Dec 4, 2024
721a9e7
Merge branch 'main' into annotate_rootOp
nithinsubbiah Dec 4, 2024
2ae4c0a
Update test name
nithinsubbiah Dec 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
static const char kConfigAttrName[] = "lowering_config";
static const char kTranslationInfoAttrName[] = "translation_info";
static const char kCompilationInfoAttrName[] = "compilation_info";
static const char kRootOpInfoAttrName[] = "root_op";

namespace mlir::iree_compiler {

Expand Down Expand Up @@ -568,4 +569,13 @@ void eraseCompilationInfo(Operation *op) {
op->removeAttr(kCompilationInfoAttrName);
}

//===----------------------------------------------------------------------===//
// Helpers for setting `iree_codegen.root_op` attribute on root operations for
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is no iree_codegen.root_op attribute. What you're doing is adding a unit attribute to the op. Perhaps rephrase it like "Helpers for setting attributes for tuner".

// tuner.
// ===----------------------------------------------------------------------===//

void setRootOpInfo(Operation *op) {
op->setAttr(kRootOpInfoAttrName, UnitAttr::get(op->getContext()));
}

} // namespace mlir::iree_compiler
5 changes: 5 additions & 0 deletions compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,12 +117,17 @@ SmallVector<Value> getTileSizes(OpBuilder &b, Operation *op, unsigned level);
/// Sets the lowering configuration, overwriting existing attribute values.
void setLoweringConfig(Operation *op, Attribute config);

/// Sets an attribute to identify the rootOp and adds any information needed for
/// the tuner from compiler. Currently, only sets a `UnitAttr`.
void setRootOpInfo(Operation *op);

/// Convenience function that sets the lowering configuration on the operation
/// and translation info.
inline LogicalResult setOpConfigAndEntryPointFnTranslation(
mlir::FunctionOpInterface entryPointFn, Operation *op,
IREE::Codegen::LoweringConfigAttrInterface config,
IREE::Codegen::TranslationInfoAttr translationInfo) {
setRootOpInfo(op);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be more specific -- If it were me, I'd add a flag in this file and wrap this statement with the flag.

if (config) {
setLoweringConfig(op, config);
}
Expand Down
1 change: 1 addition & 0 deletions compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ iree_lit_test_suite(
[
"invalid.mlir",
"lowering_config_attr.mlir",
"root_op_attribute.mlir",
"ukernel_ops.mlir",
"ukernel_ops_cse.mlir",
"workgroup_mapping_attrs.mlir",
Expand Down
1 change: 1 addition & 0 deletions compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ iree_lit_test_suite(
SRCS
"invalid.mlir"
"lowering_config_attr.mlir"
"root_op_attribute.mlir"
"ukernel_ops.mlir"
"ukernel_ops_cse.mlir"
"workgroup_mapping_attrs.mlir"
Expand Down
12 changes: 12 additions & 0 deletions compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/root_op_attribute.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
// RUN: iree-opt --split-input-file --iree-gpu-test-target=gfx942 --pass-pipeline='builtin.module(iree-llvmgpu-select-lowering-strategy)' %s | FileCheck %s
nithinsubbiah marked this conversation as resolved.
Show resolved Hide resolved

func.func @matmul(%lhs: tensor<4x4xf32>, %rhs: tensor<4x4xf32>) -> tensor<4x4xf32> {
%c0 = arith.constant 0.0 : f32
%init = tensor.empty() : tensor<4x4xf32>
%fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32>
%result = linalg.matmul ins(%lhs, %rhs: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%fill: tensor<4x4xf32>) -> tensor<4x4xf32>
return %result : tensor<4x4xf32>
}

// CHECK: %2 = linalg.matmul {lowering_config = #config, root_op} ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%1 : tensor<4x4xf32>) -> tensor<4x4xf32>
18 changes: 9 additions & 9 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1638,7 +1638,7 @@ func.func @pad_only() attributes {hal.executable.target = #executable_target_emb
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: tensor.pad {{.+}} {
// CHECK: tensor.yield
// CHECK-NEXT: } {lowering_config = #[[CONFIG]]}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

or alternatively, only match lowering_config = #[[COFIG]] in this test to allow for other attributes

// CHECK-NEXT: } {lowering_config = #[[CONFIG]], root_op}

// -----

Expand All @@ -1665,7 +1665,7 @@ func.func @winograd_output_transform() attributes {hal.executable.target = #exec
// CHECK: func.func @winograd_output_transform()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: iree_linalg_ext.winograd.output_transform
// CHECK-SAME: {lowering_config = #[[CONFIG]]}
// CHECK-SAME: {lowering_config = #[[CONFIG]], root_op}

// -----

Expand All @@ -1692,7 +1692,7 @@ func.func @winograd_input_transform() attributes {hal.executable.target = #execu
// CHECK: func.func @winograd_input_transform()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: iree_linalg_ext.winograd.input_transform
// CHECK-SAME: {lowering_config = #[[CONFIG]]}
// CHECK-SAME: {lowering_config = #[[CONFIG]], root_op}

// -----

Expand All @@ -1719,7 +1719,7 @@ func.func @winograd_filter_transform() attributes {hal.executable.target = #exec
// CHECK: func.func @winograd_filter_transform()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: iree_linalg_ext.winograd.filter_transform
// CHECK-SAME: {lowering_config = #[[CONFIG]]}
// CHECK-SAME: {lowering_config = #[[CONFIG]], root_op}

// -----

Expand Down Expand Up @@ -1762,7 +1762,7 @@ func.func @attention() attributes {hal.executable.target = #executable_target_em
// CHECK: func.func @attention()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: iree_linalg_ext.attention
// CHECK-SAME: lowering_config = #[[CONFIG]]
// CHECK-SAME: lowering_config = #[[CONFIG]], root_op

// -----

Expand Down Expand Up @@ -1801,7 +1801,7 @@ func.func @elementwise_output_transposed() attributes {hal.executable.target = #
// CHECK: func.func @elementwise_output_transposed()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
// CHECK-SAME: {lowering_config = #[[CONFIG]]}
// CHECK-SAME: {lowering_config = #[[CONFIG]], root_op}

// -----

Expand Down Expand Up @@ -1886,7 +1886,7 @@ func.func @custom_op(%arg0 : tensor<384x512xf32>, %arg1 : tensor<512x128xf32>,
// CHECK: func @custom_op(
// CHECK-SAME: translation_info = #translation
// CHECK: iree_linalg_ext.custom_op
// CHECK-SAME: attributes {lowering_config = #[[CONFIG0]]}
// CHECK-SAME: attributes {lowering_config = #[[CONFIG0]], root_op}
// CHECK: ^bb
// CHECK: linalg.fill
// CHECK-SAME: {lowering_config = #[[CONFIG1]]}
Expand Down Expand Up @@ -1957,7 +1957,7 @@ func.func @test_tiling_cpu_default(%arg0: tensor<256x256xi8>, %arg1: tensor<256x
// CHECK-DAG: #[[TRANSLATION_INFO]] = #iree_codegen.translation_info<pipeline = CPUDefault>
// CHECK: func @test_tiling_cpu_default(
// CHECK-SAME: translation_info = #[[TRANSLATION_INFO]]
// CHECK: linalg.quantized_matmul {lowering_config = #[[CONFIG0]]}
// CHECK: linalg.quantized_matmul {lowering_config = #[[CONFIG0]], root_op}

// -----

Expand All @@ -1982,7 +1982,7 @@ func.func @i1_type() attributes {hal.executable.target = #executable_target_emb
// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[8], [8], [0], [0]]>
// CHECK: func @i1_type()
// CHECK: linalg.generic {
// CHECK-SAME: {lowering_config = #[[CONFIG]]}
// CHECK-SAME: {lowering_config = #[[CONFIG]], root_op}

// -----
#pipeline_layout = #hal.pipeline.layout<bindings = [
Expand Down
Loading