derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
42 lines
2.6 KiB
MLIR
42 lines
2.6 KiB
MLIR
// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s
|
|
|
|
// CHECK-LABEL: masked_matmul
|
|
func.func @masked_matmul(%module: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
|
|
|
|
// CHECK: %[[MLHS:.*]] = vector.create_mask {{.*}} : vector<8x8xi1>
|
|
// CHECK: %[[LHS:.*]] = vector.transfer_read %{{.*}}, %[[MLHS]] {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, 1], offset: ?>>, vector<8x8xf32>
|
|
// CHECK: %[[MRHS:.*]] = vector.create_mask {{.*}} : vector<8x8xi1>
|
|
// CHECK: %[[RHS:.*]] = vector.transfer_read %{{.*}}, %[[MRHS]] {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, 1], offset: ?>>, vector<8x8xf32>
|
|
// CHECK: %[[MACC:.*]] = vector.create_mask {{.*}} : vector<8x8xi1>
|
|
// CHECK: %[[ACC:.*]] = vector.transfer_read {{.*}}, %[[MACC]] {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, 1], offset: ?>>, vector<8x8xf32>
|
|
// CHECK: %[[MRES:.*]] = vector.create_mask {{.*}} : vector<8x8x8xi1>
|
|
// CHECK: %[[RES:.*]] = vector.mask %[[MRES]] { vector.contract
|
|
// CHECK-SAME: : vector<8x8xf32>, vector<8x8xf32> into vector<8x8xf32>
|
|
// CHECK-SAME: : vector<8x8x8xi1> -> vector<8x8xf32>
|
|
// CHECK: vector.transfer_write %[[RES]], %{{.*}}, %[[MACC]] {in_bounds = [true, true]} : vector<8x8xf32>, memref<?x?xf32, strided<[?, 1], offset: ?>>
|
|
linalg.matmul ins(%module, %arg1 : memref<?x?xf32>, memref<?x?xf32>) outs(%arg2 : memref<?x?xf32>)
|
|
return
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["linalg.matmul"]} in %module
|
|
: (!transform.any_op) -> !transform.any_op
|
|
%tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [64, 128, 256]
|
|
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
|
|
%tiled_linalg_op_0, %loops_1:3 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [8, 8, 8]
|
|
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
|
|
transform.structured.vectorize %tiled_linalg_op_0 vector_sizes [8, 8, 8]
|
|
: !transform.any_op
|
|
|
|
%func = transform.structured.match ops{["func.func"]} in %module
|
|
: (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.vector.lower_masked_transfers
|
|
transform.apply_patterns.vector.transfer_permutation_patterns
|
|
transform.apply_patterns.vector.reduction_to_contract
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|