Skip to content

Commit

Permalink
Attaching pipeline layout to hal.interface.binding.subspan & co. (#18098
Browse files Browse the repository at this point in the history 
)

This allows for the whole layout to be known locally when lowering out
of HAL and into target-specific binding data structures. This
information was (and is still) available on the exports but annoying to
get to and not present in all tests. This allowed removing the
descriptor type from the subspan op and will allow for us to have
non-i32 push constant types in the future. Verifiers were added for both
push constant and descriptor set/binding ordinals now that the
information is cheap to verify.

Progress on #17875 (this is needed for lowering non-0 ordinal descriptor
sets to CPU/CUDA/ROCM targets).
  • Loading branch information
@benvanik
benvanik committed Aug 5, 2024
1 parent 5ca6bee commit 2193406
Show file tree
Hide file tree
Showing 204 changed files with 14,760 additions and 10,766 deletions.
881 changes: 544 additions & 337 deletions compiler/src/iree/compiler/Codegen/Common/CPU/test/llvmcpu_materialize_encoding.mlir
View file

Large diffs are not rendered by default.

125 changes: 79 additions & 46 deletions compiler/src/iree/compiler/Codegen/Common/CPU/test/vmvx_materialize_encoding.mlir
View file

Large diffs are not rendered by default.

4 changes: 2 additions & 2 deletions compiler/src/iree/compiler/Codegen/Common/ConvertBf16ToUInt16Buffers.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,8 +89,8 @@ struct ConvertHalInterfaceBindingSubspan final

auto newOp =
rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceBindingSubspanOp>(
op, newResultTy, adaptor.getSet(), adaptor.getBinding(),
adaptor.getDescriptorType(), adaptor.getByteOffset(),
op, newResultTy, adaptor.getLayout(), adaptor.getSet(),
adaptor.getBinding(), adaptor.getByteOffset(),
adaptor.getDynamicDims(), adaptor.getAlignmentAttr(),
adaptor.getDescriptorFlagsAttr());
LLVM_DEBUG(llvm::dbgs() << "Bf16Emulation: new op: " << newOp << "\n");
Expand Down
4 changes: 2 additions & 2 deletions compiler/src/iree/compiler/Codegen/Common/EmulateNarrowType.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,8 +78,8 @@ struct ConvertHalInterfaceBindingSubspan final
}

rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceBindingSubspanOp>(
op, newResultType, adaptor.getSet(), adaptor.getBinding(),
adaptor.getDescriptorType(), byteOffset, dynamicLinearizedSize,
op, newResultType, adaptor.getLayout(), adaptor.getSet(),
adaptor.getBinding(), byteOffset, dynamicLinearizedSize,
adaptor.getAlignmentAttr(), adaptor.getDescriptorFlagsAttr());
return success();
}
Expand Down
4 changes: 2 additions & 2 deletions compiler/src/iree/compiler/Codegen/Common/FlattenMemRefSubspanPass.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -283,8 +283,8 @@ struct FlattenBindingSubspan final

auto newOffset = rewriter.create<arith::ConstantIndexOp>(loc, 0);
auto newOp = rewriter.create<IREE::HAL::InterfaceBindingSubspanOp>(
subspanOp.getLoc(), newType, subspanOp.getSet(), subspanOp.getBinding(),
subspanOp.getDescriptorType(), newOffset, dynamicShape,
subspanOp.getLoc(), newType, subspanOp.getLayout(), subspanOp.getSet(),
subspanOp.getBinding(), newOffset, dynamicShape,
subspanOp.getAlignmentAttr(), subspanOp.getDescriptorFlagsAttr());

Value replacement = newOp;
Expand Down
259 changes: 153 additions & 106 deletions compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
View file

Large diffs are not rendered by default.

16 changes: 12 additions & 4 deletions compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_create_fast_slow_path.mlir
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,22 @@
// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-create-fast-slow-path))" --mlir-print-local-scope %s | FileCheck %s

#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
#hal.descriptor_set.binding<1, storage_buffer>,
#hal.descriptor_set.binding<2, storage_buffer>,
#hal.descriptor_set.binding<3, storage_buffer>
]>
]>
func.func @padded_conv() {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c32 = arith.constant 32 : index
%c112 = arith.constant 112 : index
%0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<1x224x224x3xf32>>
%1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<3x3x3x32xf32>>
%2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<1x112x112x32xf32>>
%3 = hal.interface.binding.subspan set(0) binding(3) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<1x112x112x32xf32>>
%0 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<1x224x224x3xf32>>
%1 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<3x3x3x32xf32>>
%2 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(2) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<1x112x112x32xf32>>
%3 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(3) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<1x112x112x32xf32>>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_count_x = hal.interface.workgroup.count[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
Expand Down
118 changes: 64 additions & 54 deletions compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_distribute.mlir
View file
Original file line number Diff line number Diff line change
@@ -1,36 +1,41 @@
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-distribute, cse))" %s --split-input-file | FileCheck %s

#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
#hal.descriptor_set.binding<1, storage_buffer>,
#hal.descriptor_set.binding<2, storage_buffer>
]>
]>
#map = affine_map<()[s0] -> (s0 * 256)>
#map1 = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>
#map2 = affine_map<(d0) -> (d0 * 4)>
#translation = #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [64, 1, 1]>
module {
func.func @add_tensor() attributes {translation_info = #translation} {
%cst = arith.constant 0.000000e+00 : f32
%c64 = arith.constant 64 : index
%c0 = arith.constant 0 : index
%0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %0, 64 : memref<233x1024xf32>
%1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %1, 64 : memref<233x1024xf32>
%2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %2, 64 : memref<233x1024xf32>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
%3 = affine.apply #map()[%workgroup_id_x]
%subview = memref.subview %2[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_0 = memref.subview %0[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_1 = memref.subview %1[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
scf.forall (%arg0) in (%c64) {
%4 = affine.apply #map2(%arg0)
%subview_2 = memref.subview %subview[0, %4] [1, 4] [1, 1] : memref<1x256xf32, #map1> to memref<1x4xf32, #map1>
%5 = vector.transfer_read %subview_0[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%6 = vector.transfer_read %subview_1[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%7 = arith.addf %5, %6 : vector<4xf32>
vector.transfer_write %7, %subview_2[%c0, %c0] {in_bounds = [true]} : vector<4xf32>, memref<1x4xf32, #map1>
} {mapping = [#gpu.thread<x>]}
return
}
func.func @add_tensor() attributes {translation_info = #translation} {
%cst = arith.constant 0.000000e+00 : f32
%c64 = arith.constant 64 : index
%c0 = arith.constant 0 : index
%0 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %0, 64 : memref<233x1024xf32>
%1 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %1, 64 : memref<233x1024xf32>
%2 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(2) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %2, 64 : memref<233x1024xf32>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
%3 = affine.apply #map()[%workgroup_id_x]
%subview = memref.subview %2[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_0 = memref.subview %0[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_1 = memref.subview %1[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
scf.forall (%arg0) in (%c64) {
%4 = affine.apply #map2(%arg0)
%subview_2 = memref.subview %subview[0, %4] [1, 4] [1, 1] : memref<1x256xf32, #map1> to memref<1x4xf32, #map1>
%5 = vector.transfer_read %subview_0[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%6 = vector.transfer_read %subview_1[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%7 = arith.addf %5, %6 : vector<4xf32>
vector.transfer_write %7, %subview_2[%c0, %c0] {in_bounds = [true]} : vector<4xf32>, memref<1x4xf32, #map1>
} {mapping = [#gpu.thread<x>]}
return
}

// CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0 * 4)>
Expand All @@ -46,37 +51,42 @@ module {

// -----

#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
#hal.descriptor_set.binding<1, storage_buffer>,
#hal.descriptor_set.binding<2, storage_buffer>
]>
]>
#map = affine_map<()[s0] -> (s0 * 256)>
#map1 = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>
#map2 = affine_map<(d0) -> (d0 * 4)>
#translation = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [64, 1, 1]>
module {
func.func @add_tensor_lane_id() attributes {translation_info = #translation} {
%cst = arith.constant 0.000000e+00 : f32
%c64 = arith.constant 64 : index
%c0 = arith.constant 0 : index
%0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %0, 64 : memref<233x1024xf32>
%1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %1, 64 : memref<233x1024xf32>
%2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %2, 64 : memref<233x1024xf32>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
%3 = affine.apply #map()[%workgroup_id_x]
%subview = memref.subview %2[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_0 = memref.subview %0[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_1 = memref.subview %1[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
scf.forall (%arg0) in (%c64) {
%4 = affine.apply #map2(%arg0)
%subview_2 = memref.subview %subview[0, %4] [1, 4] [1, 1] : memref<1x256xf32, #map1> to memref<1x4xf32, #map1>
%5 = vector.transfer_read %subview_0[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%6 = vector.transfer_read %subview_1[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%7 = arith.addf %5, %6 : vector<4xf32>
vector.transfer_write %7, %subview_2[%c0, %c0] {in_bounds = [true]} : vector<4xf32>, memref<1x4xf32, #map1>
} {mapping = [#iree_gpu.lane_id<0>]}
return
}
func.func @add_tensor_lane_id() attributes {translation_info = #translation} {
%cst = arith.constant 0.000000e+00 : f32
%c64 = arith.constant 64 : index
%c0 = arith.constant 0 : index
%0 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %0, 64 : memref<233x1024xf32>
%1 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %1, 64 : memref<233x1024xf32>
%2 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(2) alignment(64) offset(%c0) : memref<233x1024xf32>
memref.assume_alignment %2, 64 : memref<233x1024xf32>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
%3 = affine.apply #map()[%workgroup_id_x]
%subview = memref.subview %2[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_0 = memref.subview %0[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
%subview_1 = memref.subview %1[%workgroup_id_y, %3] [1, 256] [1, 1] : memref<233x1024xf32> to memref<1x256xf32, #map1>
scf.forall (%arg0) in (%c64) {
%4 = affine.apply #map2(%arg0)
%subview_2 = memref.subview %subview[0, %4] [1, 4] [1, 1] : memref<1x256xf32, #map1> to memref<1x4xf32, #map1>
%5 = vector.transfer_read %subview_0[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%6 = vector.transfer_read %subview_1[%c0, %4], %cst {in_bounds = [true]} : memref<1x256xf32, #map1>, vector<4xf32>
%7 = arith.addf %5, %6 : vector<4xf32>
vector.transfer_write %7, %subview_2[%c0, %c0] {in_bounds = [true]} : vector<4xf32>, memref<1x4xf32, #map1>
} {mapping = [#iree_gpu.lane_id<0>]}
return
}

// CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0 * 4)>
Expand Down
43 changes: 31 additions & 12 deletions compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_pipeline.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,13 @@
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-pipelining{epilogue-peeling=false}))" --split-input-file %s | FileCheck %s
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-pipelining{pipeline-depth=3 schedule-index=2 epilogue-peeling=false}))" --split-input-file %s | FileCheck -check-prefix=CHECK-NV %s


#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
#hal.descriptor_set.binding<1, storage_buffer>,
#hal.descriptor_set.binding<2, storage_buffer>
]>
]>
func.func @_matmul_f16_f16_dispatch_0_fill_3456x1024() {
%c2048 = arith.constant 2048 : index
%c32 = arith.constant 32 : index
Expand All @@ -15,11 +21,11 @@ func.func @_matmul_f16_f16_dispatch_0_fill_3456x1024() {
%3 = gpu.thread_id z
%4 = memref.alloc() : memref<4x32x40xf16, 3>
%5 = memref.alloc() : memref<4x32x40xf16, 3>
%6 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : memref<3456x2048xf16>
%6 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) : memref<3456x2048xf16>
memref.assume_alignment %6, 64 : memref<3456x2048xf16>
%7 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : memref<2048x1024xf16>
%7 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) : memref<2048x1024xf16>
memref.assume_alignment %7, 64 : memref<2048x1024xf16>
%8 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : memref<3456x1024xf16>
%8 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(2) alignment(64) offset(%c0) : memref<3456x1024xf16>
memref.assume_alignment %8, 64 : memref<3456x1024xf16>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
Expand Down Expand Up @@ -57,6 +63,13 @@ func.func @_matmul_f16_f16_dispatch_0_fill_3456x1024() {

// -----

#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
#hal.descriptor_set.binding<1, storage_buffer>,
#hal.descriptor_set.binding<2, storage_buffer>
]>
]>
func.func @nvidia_tenscore_schedule_f16() {
%c3 = arith.constant 3 : index
%c31 = arith.constant 31 : index
Expand All @@ -73,11 +86,11 @@ func.func @nvidia_tenscore_schedule_f16() {
%alloc = memref.alloc() : memref<128x256xf16, #gpu.address_space<workgroup>>
%alloc_1 = memref.alloc() : memref<3x128x32xf16, #gpu.address_space<workgroup>>
%alloc_2 = memref.alloc() : memref<3x32x256xf16, #gpu.address_space<workgroup>>
%3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<512x1280xf16>
%3 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : memref<512x1280xf16>
memref.assume_alignment %3, 64 : memref<512x1280xf16>
%4 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<1280x1280xf16>
%4 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : memref<1280x1280xf16>
memref.assume_alignment %4, 64 : memref<1280x1280xf16>
%5 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : memref<512x1280xf16>
%5 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(2) alignment(64) offset(%c0) : memref<512x1280xf16>
memref.assume_alignment %5, 64 : memref<512x1280xf16>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
Expand Down Expand Up @@ -503,6 +516,14 @@ func.func @nvidia_tenscore_schedule_f16() {
// CHECK-NV: vector.store

// -----

#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
#hal.descriptor_set.binding<1, storage_buffer>,
#hal.descriptor_set.binding<2, storage_buffer>
]>
]>
func.func @nvidia_tenscore_schedule_f32() {
%c31 = arith.constant 31 : index
%c2 = arith.constant 2 : index
Expand All @@ -519,11 +540,11 @@ func.func @nvidia_tenscore_schedule_f32() {
%alloc = memref.alloc() : memref<128x128xf32, #gpu.address_space<workgroup>>
%alloc_2 = memref.alloc() : memref<3x128x32xf32, #gpu.address_space<workgroup>>
%alloc_3 = memref.alloc() : memref<3x32x128xf32, #gpu.address_space<workgroup>>
%3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<256x256xf32>
%3 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : memref<256x256xf32>
memref.assume_alignment %3, 64 : memref<256x256xf32>
%4 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<256x256xf32>
%4 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : memref<256x256xf32>
memref.assume_alignment %4, 64 : memref<256x256xf32>
%5 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : memref<256x256xf32>
%5 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(2) alignment(64) offset(%c0) : memref<256x256xf32>
memref.assume_alignment %5, 64 : memref<256x256xf32>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
Expand Down Expand Up @@ -1345,5 +1366,3 @@ func.func @nvidia_tenscore_schedule_f32() {
// CHECK-NV-COUNT-32: nvgpu.mma.sync
// CHECK-NV: }
// CHECK-NV: vector.store

// -----
Loading

0 comments on commit 2193406

Please sign in to comment.