[Codegen][CPU] Add FP32 packed layouts for Wasm SIMD (#25199) While compiling FP32 matmuls to Wasm, I noticed that enabling data tiling wasn't selecting a packed layout because the CPU layout selector didn't have a Wasm case. This adds tile choices for Wasm SIMD128. I've added tests for wasm32/wasm64, including checks that scalar Wasm and FP64 keep their existing behavior. --------- Signed-off-by: anutosh491 <andersonbhat491@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel index 623eaf3..f81d823 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
@@ -105,6 +105,7 @@ "materialize_encoding_riscv.mlir", "materialize_encoding_riscv_inner_tiled.mlir", "materialize_encoding_vmvx.mlir", + "materialize_encoding_wasm.mlir", "materialize_encoding_x86_64.mlir", "materialize_smt_constraint_assignments.mlir", "materialize_tuning_specs.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt index b0c5465..a2ebfdf 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
@@ -100,6 +100,7 @@ "materialize_encoding_riscv.mlir" "materialize_encoding_riscv_inner_tiled.mlir" "materialize_encoding_vmvx.mlir" + "materialize_encoding_wasm.mlir" "materialize_encoding_x86_64.mlir" "materialize_smt_constraint_assignments.mlir" "materialize_tuning_specs.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_wasm.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_wasm.mlir new file mode 100644 index 0000000..c3e21c2 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_wasm.mlir
@@ -0,0 +1,151 @@ +// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --split-input-file %s | FileCheck %s + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +func.func @simd32(%lhs: tensor<?x?xf32, #lhs>, %rhs: tensor<?x?xf32, #rhs>, %acc: tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_32", {target_triple = "wasm32-unknown-emscripten", cpu_features = "+simd128", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf32, #lhs>, tensor<?x?xf32, #rhs>) outs(%acc : tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> + return %result : tensor<?x?xf32, #out> +} +// CHECK-LABEL: func @simd32( +// CHECK-SAME: tensor<?x?x4x1xf32> +// CHECK-SAME: tensor<?x?x8x1xf32> +// CHECK-SAME: tensor<?x?x4x8xf32> +// CHECK: linalg.mmt4d + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [2, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [2, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [2, ?, ?]> +func.func @simd32_narrow_m2(%lhs: tensor<2x?xf32, #lhs>, %rhs: tensor<?x?xf32, #rhs>, %acc: tensor<2x?xf32, #out>) -> tensor<2x?xf32, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_32", {target_triple = "wasm32-unknown-emscripten", cpu_features = "+simd128", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<2x?xf32, #lhs>, tensor<?x?xf32, #rhs>) outs(%acc : tensor<2x?xf32, #out>) -> tensor<2x?xf32, #out> + return %result : tensor<2x?xf32, #out> +} +// CHECK-LABEL: func @simd32_narrow_m2( +// CHECK-SAME: tensor<1x?x2x1xf32> +// CHECK-SAME: tensor<?x?x8x1xf32> +// CHECK-SAME: tensor<1x?x2x8xf32> +// CHECK: linalg.mmt4d + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [1, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [1, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [1, ?, ?]> +func.func @simd32_narrow_m1(%lhs: tensor<1x?xf32, #lhs>, %rhs: tensor<?x?xf32, #rhs>, %acc: tensor<1x?xf32, #out>) -> tensor<1x?xf32, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_32", {target_triple = "wasm32-unknown-emscripten", cpu_features = "+simd128", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<1x?xf32, #lhs>, tensor<?x?xf32, #rhs>) outs(%acc : tensor<1x?xf32, #out>) -> tensor<1x?xf32, #out> + return %result : tensor<1x?xf32, #out> +} +// CHECK-LABEL: func @simd32_narrow_m1( +// CHECK-SAME: tensor<1x?x1x1xf32> +// CHECK-SAME: tensor<?x?x8x1xf32> +// CHECK-SAME: tensor<1x?x1x8xf32> +// CHECK: linalg.mmt4d + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +func.func @simd64(%lhs: tensor<?x?xf32, #lhs>, %rhs: tensor<?x?xf32, #rhs>, %acc: tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_64", {target_triple = "wasm64-unknown-emscripten", cpu_features = "+simd128", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf32, #lhs>, tensor<?x?xf32, #rhs>) outs(%acc : tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> + return %result : tensor<?x?xf32, #out> +} +// CHECK-LABEL: func @simd64( +// CHECK-SAME: tensor<?x?x4x1xf32> +// CHECK-SAME: tensor<?x?x8x1xf32> +// CHECK-SAME: tensor<?x?x4x8xf32> +// CHECK: linalg.mmt4d + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +func.func @scalar32(%lhs: tensor<?x?xf32, #lhs>, %rhs: tensor<?x?xf32, #rhs>, %acc: tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_32", {target_triple = "wasm32-unknown-emscripten", cpu_features = "", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf32, #lhs>, tensor<?x?xf32, #rhs>) outs(%acc : tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> + return %result : tensor<?x?xf32, #out> +} +// CHECK-LABEL: func @scalar32( +// CHECK-NOT: linalg.mmt4d +// CHECK: linalg.matmul + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +func.func @scalar64(%lhs: tensor<?x?xf32, #lhs>, %rhs: tensor<?x?xf32, #rhs>, %acc: tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_64", {target_triple = "wasm64-unknown-emscripten", cpu_features = "", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf32, #lhs>, tensor<?x?xf32, #rhs>) outs(%acc : tensor<?x?xf32, #out>) -> tensor<?x?xf32, #out> + return %result : tensor<?x?xf32, #out> +} +// CHECK-LABEL: func @scalar64( +// CHECK-NOT: linalg.mmt4d +// CHECK: linalg.matmul + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f64, f64, f64], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f64, f64, f64], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f64, f64, f64], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +func.func @f64_32(%lhs: tensor<?x?xf64, #lhs>, %rhs: tensor<?x?xf64, #rhs>, %acc: tensor<?x?xf64, #out>) -> tensor<?x?xf64, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_32", {target_triple = "wasm32-unknown-emscripten", cpu_features = "+simd128", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf64, #lhs>, tensor<?x?xf64, #rhs>) outs(%acc : tensor<?x?xf64, #out>) -> tensor<?x?xf64, #out> + return %result : tensor<?x?xf64, #out> +} +// CHECK-LABEL: func @f64_32( +// CHECK-NOT: linalg.mmt4d +// CHECK: linalg.matmul + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f64, f64, f64], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f64, f64, f64], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +#out = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f64, f64, f64], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]> +func.func @f64_64(%lhs: tensor<?x?xf64, #lhs>, %rhs: tensor<?x?xf64, #rhs>, %acc: tensor<?x?xf64, #out>) -> tensor<?x?xf64, #out> attributes { + hal.executable.target = #hal.executable.target<"llvm-cpu", "system-wasm-wasm_64", {target_triple = "wasm64-unknown-emscripten", cpu_features = "+simd128", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}> +} { + %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf64, #lhs>, tensor<?x?xf64, #rhs>) outs(%acc : tensor<?x?xf64, #out>) -> tensor<?x?xf64, #out> + return %result : tensor<?x?xf64, #out> +} +// CHECK-LABEL: func @f64_64( +// CHECK-NOT: linalg.mmt4d +// CHECK: linalg.matmul
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp index b38c6d4..d8c214c 100644 --- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp +++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -1558,11 +1558,35 @@ return {}; } +// Enumerate tile sizes to choose from on Wasm. +// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases +// are handled by transposition in chooseMatmulTile. +static SmallVector<TileMxNxK> enumerateMatmulTileWasm(TypeRange elementTypes, + DictionaryAttr config) { + if (hasFeature(config, "+simd128") && + llvm::all_of(elementTypes, [](Type type) { return type.isF32(); })) { + // SIMD128 has four FP32 lanes: N0=8 exposes two SIMD chunks per row, + // and M0=4 exposes four rows for RHS reuse and LHS scalar broadcasts. + // The 4x8 tile needs eight v128 accumulator values, leaving headroom for + // RHS vectors, broadcasts and mul/add temporaries. This is a live-value + // budget: the Wasm engine allocates physical registers for the host. + // K0=1 requests no extra inner-K unrolling through the packed layout. + // The narrow-M tiles reduce accumulator values to four/two and avoid + // padding one/two-row matmuls to four rows. + return {TileMxNxK{4, 8, 1}, TileMxNxK{2, 8, 1}, TileMxNxK{1, 8, 1}}; + } + // Fallback - no architecture-optimized tile size for this case. + return {}; +} + static SmallVector<TileMxNxK> enumerateCPUMatmulTiles(IREE::Encoding::EncodingAttr encoding, DictionaryAttr config) { // Enumerate available tile shapes for the given encoding and config. SmallVector<Type> elementTypes = encoding.getElementTypesArray(); + if (isWasm(config)) { + return enumerateMatmulTileWasm(elementTypes, config); + } if (isAArch64(config)) { return enumerateMatmulTileArm64(elementTypes, config); }
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp index 7446a76..9d28a2b 100644 --- a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp +++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
@@ -362,6 +362,11 @@ return triple && triple.value().isRISCV64(); } +bool isWasm(DictionaryAttr targetConfig) { + std::optional<llvm::Triple> triple = getTargetTriple(targetConfig); + return triple && triple.value().isWasm(); +} + std::array<int64_t, 3> getMaxWorkgroupCount(DictionaryAttr targetConfig) { // TODO(MaheshRavishankar): For now the target info is only available for // GPUs, and is recorded in the configuration with the name `iree.gpu.target`.
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Utils/Utils.h index 2be0923..ffa4b56 100644 --- a/compiler/src/iree/compiler/Codegen/Utils/Utils.h +++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
@@ -127,6 +127,7 @@ bool isRISCV(DictionaryAttr targetConfig); bool isRISCV32(DictionaryAttr targetConfig); bool isRISCV64(DictionaryAttr targetConfig); +bool isWasm(DictionaryAttr targetConfig); /// Get maximum workgroup count in [x, y, z] for target attribute if it is /// available. Returns ShapedType::kDynamic if it is unknown.