Skip to content

Commit

Permalink
Use workgroup_count_from_slice in Stream builtins (#18924)
Browse files Browse the repository at this point in the history 
`workgroup_count_from_dag_root` is planned to be replaced in the future
and is not supported by all codegen paths. Switch to
`workgroup_count_from_slice`.
  • Loading branch information
@qedawkins
qedawkins authored Oct 29, 2024
1 parent 67ba1c4 commit f4a5f13
Show file tree
Hide file tree
Showing 6 changed files with 30 additions and 24 deletions.
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/Builtins/fill_i16.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@

stream.executable private @__builtin_fill_i16 {
stream.executable.export public @__builtin_fill_i16 workgroups(%arg0: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg0
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @__builtin_fill_i16(%value: i16, %count: index, %out_binding: !stream.binding) {
%c0 = arith.constant 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count}
%0 = tensor.empty(%count) : tensor<?xi16>
%count0 = flow.dispatch.workload.ordinal %count, 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count0}
%0 = tensor.empty(%count0) : tensor<?xi16>
%1 = linalg.fill ins(%value : i16) outs(%0 : tensor<?xi16>) -> tensor<?xi16>
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count], strides = [1] : tensor<?xi16> -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count}
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count0], strides = [1] : tensor<?xi16> -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count}
return
}
}
Expand Down
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/Builtins/fill_i32.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@

stream.executable private @__builtin_fill_i32 {
stream.executable.export public @__builtin_fill_i32 workgroups(%arg0: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg0
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @__builtin_fill_i32(%value: i32, %count: index, %out_binding: !stream.binding) {
%c0 = arith.constant 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count}
%0 = tensor.empty(%count) : tensor<?xi32>
%count0 = flow.dispatch.workload.ordinal %count, 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count0}
%0 = tensor.empty(%count0) : tensor<?xi32>
%1 = linalg.fill ins(%value : i32) outs(%0 : tensor<?xi32>) -> tensor<?xi32>
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count], strides = [1] : tensor<?xi32> -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count}
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count0], strides = [1] : tensor<?xi32> -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count}
return
}
}
Expand Down
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/Builtins/fill_i8.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@

stream.executable private @__builtin_fill_i8 {
stream.executable.export public @__builtin_fill_i8 workgroups(%arg0: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg0
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @__builtin_fill_i8(%value: i8, %count: index, %out_binding: !stream.binding) {
%c0 = arith.constant 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count}
%0 = tensor.empty(%count) : tensor<?xi8>
%count0 = flow.dispatch.workload.ordinal %count, 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count0}
%0 = tensor.empty(%count0) : tensor<?xi8>
%1 = linalg.fill ins(%value : i8) outs(%0 : tensor<?xi8>) -> tensor<?xi8>
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count], strides = [1] : tensor<?xi8> -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count}
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count0], strides = [1] : tensor<?xi8> -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count}
return
}
}
Expand Down
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/Builtins/splat_i16.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@

stream.executable private @__builtin_splat_i16 {
stream.executable.export public @__builtin_splat_i16 workgroups(%arg0: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg0
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @__builtin_splat_i16(%value: i16, %count: index, %out_binding: !stream.binding) {
%c0 = arith.constant 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count}
%0 = tensor.empty(%count) : tensor<?xi16>
%count0 = flow.dispatch.workload.ordinal %count, 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count0}
%0 = tensor.empty(%count0) : tensor<?xi16>
%1 = linalg.fill ins(%value : i16) outs(%0 : tensor<?xi16>) -> tensor<?xi16>
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count], strides = [1] : tensor<?xi16> -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count}
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count0], strides = [1] : tensor<?xi16> -> !flow.dispatch.tensor<writeonly:tensor<?xi16>>{%count}
return
}
}
Expand Down
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/Builtins/splat_i32.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@

stream.executable private @__builtin_splat_i32 {
stream.executable.export public @__builtin_splat_i32 workgroups(%arg0: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg0
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @__builtin_splat_i32(%value: i32, %count: index, %out_binding: !stream.binding) {
%c0 = arith.constant 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count}
%0 = tensor.empty(%count) : tensor<?xi32>
%count0 = flow.dispatch.workload.ordinal %count, 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count0}
%0 = tensor.empty(%count0) : tensor<?xi32>
%1 = linalg.fill ins(%value : i32) outs(%0 : tensor<?xi32>) -> tensor<?xi32>
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count], strides = [1] : tensor<?xi32> -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count}
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count0], strides = [1] : tensor<?xi32> -> !flow.dispatch.tensor<writeonly:tensor<?xi32>>{%count}
return
}
}
Expand Down
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/Builtins/splat_i8.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@

stream.executable private @__builtin_splat_i8 {
stream.executable.export public @__builtin_splat_i8 workgroups(%arg0: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg0
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @__builtin_splat_i8(%value: i8, %count: index, %out_binding: !stream.binding) {
%c0 = arith.constant 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count}
%0 = tensor.empty(%count) : tensor<?xi8>
%count0 = flow.dispatch.workload.ordinal %count, 0 : index
%out = stream.binding.subspan %out_binding[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count0}
%0 = tensor.empty(%count0) : tensor<?xi8>
%1 = linalg.fill ins(%value : i8) outs(%0 : tensor<?xi8>) -> tensor<?xi8>
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count], strides = [1] : tensor<?xi8> -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count}
flow.dispatch.tensor.store %1, %out, offsets = [0], sizes = [%count0], strides = [1] : tensor<?xi8> -> !flow.dispatch.tensor<writeonly:tensor<?xi8>>{%count}
return
}
}
Expand Down

0 comments on commit f4a5f13

Please sign in to comment.