data-tiling for `f16` and `bf16` matmuls (#14207)
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp
index b854beb..c2375d8 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp
@@ -33,6 +33,15 @@
chooseMatmulTileParamsAArch64(MatmulType type, ExecutableTargetAttr target) {
switch (type) {
case MatmulType::F32F32F32:
+ case MatmulType::F16F16F32:
+ case MatmulType::F16F16F16:
+ case MatmulType::BF16BF16F32:
+ case MatmulType::BF16BF16BF16:
+ // Note: 16-bit floating point types currently use the same tile size as
+ // f32. This makes sense when either (1) the accumulator is f32, or (2)
+ // the arithmetic will have to expand f16 to f32 in registers. We may
+ // reconsider when taking advantage of native f16/bf16 arithmetic when the
+ // accumulator itself is f16/bf16.
return {8, 1, 8};
case MatmulType::I8I8I32:
if (hasFeature(target, "+i8mm")) {
@@ -54,8 +63,18 @@
chooseMatmulTileParamsX86_64(MatmulType type, ExecutableTargetAttr target) {
switch (type) {
case MatmulType::F32F32F32:
- if (hasAVX512fFeature(target))
+ case MatmulType::F16F16F32:
+ case MatmulType::F16F16F16:
+ case MatmulType::BF16BF16F32:
+ case MatmulType::BF16BF16BF16:
+ // Note: 16-bit floating point types currently use the same tile size as
+ // f32. This makes sense when either (1) the accumulator is f32, or (2)
+ // the arithmetic will have to expand f16 to f32 in registers. We may
+ // reconsider when taking advantage of native f16/bf16 arithmetic when the
+ // accumulator itself is f16/bf16.
+ if (hasFeature(target, "+avx512f")) {
return {16, 1, 16};
+ }
if (hasFeature(target, "+avx")) {
// Note: for good performance, most +avx users will also want to add
// +fma, but that's a local instruction selection detail and the tile
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/materialize_encoding.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/materialize_encoding.mlir
index 15c7430..dc7ad00 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/materialize_encoding.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/materialize_encoding.mlir
@@ -212,6 +212,66 @@
// -----
+func.func @matmul_lowering_f16f16f16_aarch64() attributes {
+ hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="aarch64-xyz-xyz"}>
+} {
+ %c0 = arith.constant 0 : index
+ %M = hal.interface.constant.load[0] : index
+ %N = hal.interface.constant.load[1] : index
+ %K = hal.interface.constant.load[2] : index
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>>{%M, %K}
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>>{%K, %N}
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>>{%M, %N}
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [%M, %K], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>>{%M, %K}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [%K, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>>{%K, %N}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>
+ %5 = flow.dispatch.tensor.load %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>>{%M, %N}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+ %6 = linalg.matmul
+ ins(%3, %4 : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>,
+ tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>)
+ outs(%5 : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>)
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+ -> !flow.dispatch.tensor<readwrite:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>>{%M, %N}
+ return
+}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 8)>
+// CHECK: func @matmul_lowering_f16f16f16_aarch64()
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load[0]
+// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load[1]
+// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load[2]
+// CHECK-DAG: %[[TILED_M:.+]] = affine.apply #[[MAP0]]()[%[[M]]]
+// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(0)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x8x1xf16>>{%[[TILED_M]], %[[K]]}
+// CHECK: %[[TILED_N:.+]] = affine.apply #[[MAP0]]()[%[[N]]]
+// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(1)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x8x1xf16>>{%[[TILED_N]], %[[K]]}
+// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(2)
+// CHECK-SAME: !flow.dispatch.tensor<readwrite:tensor<?x?x8x8xf16>>{%[[TILED_M]], %[[TILED_N]]}
+// CHECK: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[K]], 8, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_N]], %[[K]], 8, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[OUTS:.+]] = flow.dispatch.tensor.load %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 8, 8], strides = [1, 1, 1, 1]
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: flow.dispatch.tensor.store %[[MMT4D]], %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 8, 8], strides = [1, 1, 1, 1]
+
+// -----
+
func.func @matmul_lowering_f32f32f32_x86_64() attributes {
hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="x86_64-xyz-xyz"}>
} {
@@ -393,6 +453,246 @@
// -----
+func.func @matmul_lowering_f16f16f32_x86_64_avx512f() attributes {
+ hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
+} {
+ %c0 = arith.constant 0 : index
+ %M = hal.interface.constant.load[0] : index
+ %N = hal.interface.constant.load[1] : index
+ %K = hal.interface.constant.load[2] : index
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_LHS>>>{%M, %K}
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RHS>>>{%K, %N}
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>>{%M, %N}
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [%M, %K], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_LHS>>>{%M, %K}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_LHS>>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [%K, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RHS>>>{%K, %N}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RHS>>
+ %5 = flow.dispatch.tensor.load %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>>{%M, %N}
+ -> tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>
+ %6 = linalg.matmul
+ ins(%3, %4 : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_LHS>>,
+ tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RHS>>)
+ outs(%5 : tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>)
+ -> tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>
+ -> !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>>{%M, %N}
+ return
+}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
+// CHECK: func @matmul_lowering_f16f16f32_x86_64_avx512f()
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load[0]
+// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load[1]
+// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load[2]
+// CHECK-DAG: %[[TILED_M:.+]] = affine.apply #[[MAP0]]()[%[[M]]]
+// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(0)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xf16>>{%[[TILED_M]], %[[K]]}
+// CHECK: %[[TILED_N:.+]] = affine.apply #[[MAP0]]()[%[[N]]]
+// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(1)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xf16>>{%[[TILED_N]], %[[K]]}
+// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(2)
+// CHECK-SAME: !flow.dispatch.tensor<readwrite:tensor<?x?x16x16xf32>>{%[[TILED_M]], %[[TILED_N]]}
+// CHECK: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_N]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[OUTS:.+]] = flow.dispatch.tensor.load %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: flow.dispatch.tensor.store %[[MMT4D]], %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+
+// -----
+
+func.func @matmul_lowering_f16f16f16_x86_64_avx512f() attributes {
+ hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
+} {
+ %c0 = arith.constant 0 : index
+ %M = hal.interface.constant.load[0] : index
+ %N = hal.interface.constant.load[1] : index
+ %K = hal.interface.constant.load[2] : index
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>>{%M, %K}
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>>{%K, %N}
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>>{%M, %N}
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [%M, %K], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>>{%M, %K}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [%K, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>>{%K, %N}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>
+ %5 = flow.dispatch.tensor.load %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>>{%M, %N}
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+ %6 = linalg.matmul
+ ins(%3, %4 : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>,
+ tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>)
+ outs(%5 : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>)
+ -> tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+ -> !flow.dispatch.tensor<readwrite:tensor<?x?xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>>{%M, %N}
+ return
+}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
+// CHECK: func @matmul_lowering_f16f16f16_x86_64_avx512f()
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load[0]
+// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load[1]
+// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load[2]
+// CHECK-DAG: %[[TILED_M:.+]] = affine.apply #[[MAP0]]()[%[[M]]]
+// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(0)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xf16>>{%[[TILED_M]], %[[K]]}
+// CHECK: %[[TILED_N:.+]] = affine.apply #[[MAP0]]()[%[[N]]]
+// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(1)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xf16>>{%[[TILED_N]], %[[K]]}
+// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(2)
+// CHECK-SAME: !flow.dispatch.tensor<readwrite:tensor<?x?x16x16xf16>>{%[[TILED_M]], %[[TILED_N]]}
+// CHECK: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_N]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[OUTS:.+]] = flow.dispatch.tensor.load %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: flow.dispatch.tensor.store %[[MMT4D]], %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+
+// -----
+
+func.func @matmul_lowering_bf16bf16f32_x86_64_avx512f() attributes {
+ hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
+} {
+ %c0 = arith.constant 0 : index
+ %M = hal.interface.constant.load[0] : index
+ %N = hal.interface.constant.load[1] : index
+ %K = hal.interface.constant.load[2] : index
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_LHS>>>{%M, %K}
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RHS>>>{%K, %N}
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>>{%M, %N}
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [%M, %K], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_LHS>>>{%M, %K}
+ -> tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_LHS>>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [%K, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RHS>>>{%K, %N}
+ -> tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RHS>>
+ %5 = flow.dispatch.tensor.load %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>>{%M, %N}
+ -> tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>
+ %6 = linalg.matmul
+ ins(%3, %4 : tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_LHS>>,
+ tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RHS>>)
+ outs(%5 : tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>)
+ -> tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>
+ -> !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>>{%M, %N}
+ return
+}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
+// CHECK: func @matmul_lowering_bf16bf16f32_x86_64_avx512f()
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load[0]
+// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load[1]
+// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load[2]
+// CHECK-DAG: %[[TILED_M:.+]] = affine.apply #[[MAP0]]()[%[[M]]]
+// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(0)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xbf16>>{%[[TILED_M]], %[[K]]}
+// CHECK: %[[TILED_N:.+]] = affine.apply #[[MAP0]]()[%[[N]]]
+// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(1)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xbf16>>{%[[TILED_N]], %[[K]]}
+// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(2)
+// CHECK-SAME: !flow.dispatch.tensor<readwrite:tensor<?x?x16x16xf32>>{%[[TILED_M]], %[[TILED_N]]}
+// CHECK: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_N]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[OUTS:.+]] = flow.dispatch.tensor.load %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: flow.dispatch.tensor.store %[[MMT4D]], %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+
+// -----
+
+func.func @matmul_lowering_bf16bf16bf16_x86_64_avx512f() attributes {
+ hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
+} {
+ %c0 = arith.constant 0 : index
+ %M = hal.interface.constant.load[0] : index
+ %N = hal.interface.constant.load[1] : index
+ %K = hal.interface.constant.load[2] : index
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_LHS>>>{%M, %K}
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RHS>>>{%K, %N}
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0)
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>>{%M, %N}
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [%M, %K], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_LHS>>>{%M, %K}
+ -> tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_LHS>>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [%K, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readonly:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RHS>>>{%K, %N}
+ -> tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RHS>>
+ %5 = flow.dispatch.tensor.load %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : !flow.dispatch.tensor<readwrite:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>>{%M, %N}
+ -> tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>
+ %6 = linalg.matmul
+ ins(%3, %4 : tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_LHS>>,
+ tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RHS>>)
+ outs(%5 : tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>)
+ -> tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+ : tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>
+ -> !flow.dispatch.tensor<readwrite:tensor<?x?xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>>{%M, %N}
+ return
+}
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
+// CHECK: func @matmul_lowering_bf16bf16bf16_x86_64_avx512f()
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load[0]
+// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load[1]
+// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load[2]
+// CHECK-DAG: %[[TILED_M:.+]] = affine.apply #[[MAP0]]()[%[[M]]]
+// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(0)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xbf16>>{%[[TILED_M]], %[[K]]}
+// CHECK: %[[TILED_N:.+]] = affine.apply #[[MAP0]]()[%[[N]]]
+// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(1)
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x16x1xbf16>>{%[[TILED_N]], %[[K]]}
+// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(2)
+// CHECK-SAME: !flow.dispatch.tensor<readwrite:tensor<?x?x16x16xbf16>>{%[[TILED_M]], %[[TILED_N]]}
+// CHECK: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_N]], %[[K]], 16, 1], strides = [1, 1, 1, 1]
+// CHECK: %[[OUTS:.+]] = flow.dispatch.tensor.load %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: flow.dispatch.tensor.store %[[MMT4D]], %[[OUTS_BINDING]]
+// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [%[[TILED_M]], %[[TILED_N]], 16, 16], strides = [1, 1, 1, 1]
+
+// -----
+
func.func @matmul_lowering_i8i8i32_aarch64() attributes {
hal.executable.target = #hal.executable.target<"xyz", "xyz", {target_triple="aarch64-xyz-xyz"}>
} {
diff --git a/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp
index 26ae6b4..b92223f 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp
@@ -23,11 +23,26 @@
resultElementType.isSignlessInteger(32)) {
return MatmulType::I8I8I32;
}
-
if (lhsElementType.isF32() && rhsElementType.isF32() &&
resultElementType.isF32()) {
return MatmulType::F32F32F32;
}
+ if (lhsElementType.isF16() && rhsElementType.isF16() &&
+ resultElementType.isF32()) {
+ return MatmulType::F16F16F32;
+ }
+ if (lhsElementType.isF16() && rhsElementType.isF16() &&
+ resultElementType.isF16()) {
+ return MatmulType::F16F16F16;
+ }
+ if (lhsElementType.isBF16() && rhsElementType.isBF16() &&
+ resultElementType.isF32()) {
+ return MatmulType::BF16BF16F32;
+ }
+ if (lhsElementType.isBF16() && rhsElementType.isBF16() &&
+ resultElementType.isBF16()) {
+ return MatmulType::BF16BF16BF16;
+ }
return std::nullopt;
}
@@ -40,36 +55,50 @@
return encodingAttr.getEncoding().getValue();
}
+#define IREE_GETMATMULTYPE_CASE(TYPE) \
+ case TensorEncoding::MATMUL_##TYPE##_LHS: \
+ case TensorEncoding::MATMUL_##TYPE##_RHS: \
+ case TensorEncoding::MATMUL_##TYPE##_RESULT: \
+ return MatmulType::TYPE;
+
std::optional<MatmulType> getMatmulType(TensorEncoding encoding) {
switch (encoding) {
- case TensorEncoding::MATMUL_F32F32F32_LHS:
- case TensorEncoding::MATMUL_F32F32F32_RHS:
- case TensorEncoding::MATMUL_F32F32F32_RESULT:
- return MatmulType::F32F32F32;
- case TensorEncoding::MATMUL_I8I8I32_LHS:
- case TensorEncoding::MATMUL_I8I8I32_RHS:
- case TensorEncoding::MATMUL_I8I8I32_RESULT:
- return MatmulType::I8I8I32;
+ IREE_GETMATMULTYPE_CASE(F32F32F32)
+ IREE_GETMATMULTYPE_CASE(I8I8I32)
+ IREE_GETMATMULTYPE_CASE(F16F16F32)
+ IREE_GETMATMULTYPE_CASE(F16F16F16)
+ IREE_GETMATMULTYPE_CASE(BF16BF16F32)
+ IREE_GETMATMULTYPE_CASE(BF16BF16BF16)
default:
return std::nullopt;
}
}
+#undef IREE_GETMATMULTYPE_CASE
+
+#define IREE_GETMATMULOPERANDROLE_CASE(TYPE) \
+ case TensorEncoding::MATMUL_##TYPE##_LHS: \
+ return MatmulOperandRole::LHS; \
+ case TensorEncoding::MATMUL_##TYPE##_RHS: \
+ return MatmulOperandRole::RHS; \
+ case TensorEncoding::MATMUL_##TYPE##_RESULT: \
+ return MatmulOperandRole::RESULT;
+
std::optional<MatmulOperandRole> getMatmulOperandRole(TensorEncoding encoding) {
switch (encoding) {
- case TensorEncoding::MATMUL_F32F32F32_LHS:
- case TensorEncoding::MATMUL_I8I8I32_LHS:
- return MatmulOperandRole::LHS;
- case TensorEncoding::MATMUL_F32F32F32_RHS:
- case TensorEncoding::MATMUL_I8I8I32_RHS:
- return MatmulOperandRole::RHS;
- case TensorEncoding::MATMUL_F32F32F32_RESULT:
- case TensorEncoding::MATMUL_I8I8I32_RESULT:
- return MatmulOperandRole::RESULT;
+ IREE_GETMATMULOPERANDROLE_CASE(F32F32F32)
+ IREE_GETMATMULOPERANDROLE_CASE(I8I8I32)
+ IREE_GETMATMULOPERANDROLE_CASE(F16F16F32)
+ IREE_GETMATMULOPERANDROLE_CASE(F16F16F16)
+ IREE_GETMATMULOPERANDROLE_CASE(BF16BF16F32)
+ IREE_GETMATMULOPERANDROLE_CASE(BF16BF16BF16)
+
default:
return std::nullopt;
}
}
+#undef IREE_GETMATMULOPERANDROLE_CASE
+
} // namespace iree_compiler
} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h
index 1eb914b..22cc758 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h
@@ -18,6 +18,10 @@
enum class MatmulType {
F32F32F32,
I8I8I32,
+ F16F16F32,
+ F16F16F16,
+ BF16BF16F32,
+ BF16BF16BF16
};
// Enumeration of the operands of a matmul-like operation such as linalg.matmul.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/SetEncoding.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/SetEncoding.cpp
index f9808ae..fb8fc94 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/SetEncoding.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/SetEncoding.cpp
@@ -160,6 +160,26 @@
lhsEncoding = TensorEncoding::MATMUL_F32F32F32_LHS;
rhsEncoding = TensorEncoding::MATMUL_F32F32F32_RHS;
outEncoding = TensorEncoding::MATMUL_F32F32F32_RESULT;
+ } else if (lhsElemType.isF16() && rhsElemType.isF16() &&
+ outElemType.isF32()) {
+ lhsEncoding = TensorEncoding::MATMUL_F16F16F32_LHS;
+ rhsEncoding = TensorEncoding::MATMUL_F16F16F32_RHS;
+ outEncoding = TensorEncoding::MATMUL_F16F16F32_RESULT;
+ } else if (lhsElemType.isF16() && rhsElemType.isF16() &&
+ outElemType.isF16()) {
+ lhsEncoding = TensorEncoding::MATMUL_F16F16F16_LHS;
+ rhsEncoding = TensorEncoding::MATMUL_F16F16F16_RHS;
+ outEncoding = TensorEncoding::MATMUL_F16F16F16_RESULT;
+ } else if (lhsElemType.isBF16() && rhsElemType.isBF16() &&
+ outElemType.isF32()) {
+ lhsEncoding = TensorEncoding::MATMUL_BF16BF16F32_LHS;
+ rhsEncoding = TensorEncoding::MATMUL_BF16BF16F32_RHS;
+ outEncoding = TensorEncoding::MATMUL_BF16BF16F32_RESULT;
+ } else if (lhsElemType.isBF16() && rhsElemType.isBF16() &&
+ outElemType.isBF16()) {
+ lhsEncoding = TensorEncoding::MATMUL_BF16BF16BF16_LHS;
+ rhsEncoding = TensorEncoding::MATMUL_BF16BF16BF16_RHS;
+ outEncoding = TensorEncoding::MATMUL_BF16BF16BF16_RESULT;
} else if (lhsElemType.isSignlessInteger(8) &&
rhsElemType.isSignlessInteger(8) &&
outElemType.isSignlessInteger(32)) {
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/set_encoding.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/set_encoding.mlir
index ecde468..92de5a8 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/set_encoding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/set_encoding.mlir
@@ -74,7 +74,6 @@
// PADDING: %[[RESULT:.+]] = iree_linalg_ext.unset_encoding %[[MATMUL]]
// PADDING: return %[[RESULT]]
-
// -----
func.func @matmul_dynamic(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
@@ -120,6 +119,126 @@
// -----
+func.func @matmul_i8i8i32(%arg0 : tensor<128x256xi8>, %arg1 : tensor<256x512xi8>,
+ %arg2 : tensor<128x512xi32>) -> tensor<128x512xi32> {
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<128x256xi8>, tensor<256x512xi8>)
+ outs(%arg2 : tensor<128x512xi32>) -> tensor<128x512xi32>
+ return %0 : tensor<128x512xi32>
+}
+// CHECK: func @matmul_i8i8i32(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xi8>
+// CHECK-SAME: %[[ARG1:.+]]: tensor<256x512xi8>
+// CHECK-SAME: %[[ARG2:.+]]: tensor<128x512xi32>
+// CHECK: %[[LHS:.+]] = iree_linalg_ext.set_encoding %[[ARG0]]
+// CHECK-SAME: tensor<128x256xi8, #iree_linalg_ext.encoding<MATMUL_I8I8I32_LHS>>
+// CHECK: %[[RHS:.+]] = iree_linalg_ext.set_encoding %[[ARG1]]
+// CHECK-SAME: tensor<256x512xi8, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RHS>>
+// CHECK: %[[OUTS:.+]] = iree_linalg_ext.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<128x512xi32, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RESULT>>
+// CHECK: %[[MATMUL:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: %[[RESULT:.+]] = iree_linalg_ext.unset_encoding %[[MATMUL]]
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @matmul_f16f16f32(%arg0 : tensor<128x256xf16>, %arg1 : tensor<256x512xf16>,
+ %arg2 : tensor<128x512xf32>) -> tensor<128x512xf32> {
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<128x256xf16>, tensor<256x512xf16>)
+ outs(%arg2 : tensor<128x512xf32>) -> tensor<128x512xf32>
+ return %0 : tensor<128x512xf32>
+}
+// CHECK: func @matmul_f16f16f32(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xf16>
+// CHECK-SAME: %[[ARG1:.+]]: tensor<256x512xf16>
+// CHECK-SAME: %[[ARG2:.+]]: tensor<128x512xf32>
+// CHECK: %[[LHS:.+]] = iree_linalg_ext.set_encoding %[[ARG0]]
+// CHECK-SAME: tensor<128x256xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_LHS>>
+// CHECK: %[[RHS:.+]] = iree_linalg_ext.set_encoding %[[ARG1]]
+// CHECK-SAME: tensor<256x512xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RHS>>
+// CHECK: %[[OUTS:.+]] = iree_linalg_ext.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<128x512xf32, #iree_linalg_ext.encoding<MATMUL_F16F16F32_RESULT>>
+// CHECK: %[[MATMUL:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: %[[RESULT:.+]] = iree_linalg_ext.unset_encoding %[[MATMUL]]
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @matmul_f16f16f16(%arg0 : tensor<128x256xf16>, %arg1 : tensor<256x512xf16>,
+ %arg2 : tensor<128x512xf16>) -> tensor<128x512xf16> {
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<128x256xf16>, tensor<256x512xf16>)
+ outs(%arg2 : tensor<128x512xf16>) -> tensor<128x512xf16>
+ return %0 : tensor<128x512xf16>
+}
+// CHECK: func @matmul_f16f16f16(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xf16>
+// CHECK-SAME: %[[ARG1:.+]]: tensor<256x512xf16>
+// CHECK-SAME: %[[ARG2:.+]]: tensor<128x512xf16>
+// CHECK: %[[LHS:.+]] = iree_linalg_ext.set_encoding %[[ARG0]]
+// CHECK-SAME: tensor<128x256xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_LHS>>
+// CHECK: %[[RHS:.+]] = iree_linalg_ext.set_encoding %[[ARG1]]
+// CHECK-SAME: tensor<256x512xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RHS>>
+// CHECK: %[[OUTS:.+]] = iree_linalg_ext.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<128x512xf16, #iree_linalg_ext.encoding<MATMUL_F16F16F16_RESULT>>
+// CHECK: %[[MATMUL:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: %[[RESULT:.+]] = iree_linalg_ext.unset_encoding %[[MATMUL]]
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @matmul_bf16bf16f32(%arg0 : tensor<128x256xbf16>, %arg1 : tensor<256x512xbf16>,
+ %arg2 : tensor<128x512xf32>) -> tensor<128x512xf32> {
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<128x256xbf16>, tensor<256x512xbf16>)
+ outs(%arg2 : tensor<128x512xf32>) -> tensor<128x512xf32>
+ return %0 : tensor<128x512xf32>
+}
+// CHECK: func @matmul_bf16bf16f32(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xbf16>
+// CHECK-SAME: %[[ARG1:.+]]: tensor<256x512xbf16>
+// CHECK-SAME: %[[ARG2:.+]]: tensor<128x512xf32>
+// CHECK: %[[LHS:.+]] = iree_linalg_ext.set_encoding %[[ARG0]]
+// CHECK-SAME: tensor<128x256xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_LHS>>
+// CHECK: %[[RHS:.+]] = iree_linalg_ext.set_encoding %[[ARG1]]
+// CHECK-SAME: tensor<256x512xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RHS>>
+// CHECK: %[[OUTS:.+]] = iree_linalg_ext.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<128x512xf32, #iree_linalg_ext.encoding<MATMUL_BF16BF16F32_RESULT>>
+// CHECK: %[[MATMUL:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: %[[RESULT:.+]] = iree_linalg_ext.unset_encoding %[[MATMUL]]
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @matmul_bf16bf16bf16(%arg0 : tensor<128x256xbf16>, %arg1 : tensor<256x512xbf16>,
+ %arg2 : tensor<128x512xbf16>) -> tensor<128x512xbf16> {
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<128x256xbf16>, tensor<256x512xbf16>)
+ outs(%arg2 : tensor<128x512xbf16>) -> tensor<128x512xbf16>
+ return %0 : tensor<128x512xbf16>
+}
+// CHECK: func @matmul_bf16bf16bf16(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xbf16>
+// CHECK-SAME: %[[ARG1:.+]]: tensor<256x512xbf16>
+// CHECK-SAME: %[[ARG2:.+]]: tensor<128x512xbf16>
+// CHECK: %[[LHS:.+]] = iree_linalg_ext.set_encoding %[[ARG0]]
+// CHECK-SAME: tensor<128x256xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_LHS>>
+// CHECK: %[[RHS:.+]] = iree_linalg_ext.set_encoding %[[ARG1]]
+// CHECK-SAME: tensor<256x512xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RHS>>
+// CHECK: %[[OUTS:.+]] = iree_linalg_ext.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<128x512xbf16, #iree_linalg_ext.encoding<MATMUL_BF16BF16BF16_RESULT>>
+// CHECK: %[[MATMUL:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: %[[RESULT:.+]] = iree_linalg_ext.unset_encoding %[[MATMUL]]
+// CHECK: return %[[RESULT]]
+
+// -----
+
func.func @fold_fill_with_set_encoding(%arg0 : index, %arg1 : index)
-> tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F32F32F32_LHS>> {
%cst = arith.constant 0.0 : f32
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtBase.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtBase.td
index fc228c1..68f1aec 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtBase.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtBase.td
@@ -57,12 +57,40 @@
: I32EnumAttrCase<"MATMUL_I8I8I32_RHS", 4>;
def MATMUL_I8I8I32_RESULT
: I32EnumAttrCase<"MATMUL_I8I8I32_RESULT", 5>;
+def MATMUL_F16F16F32_LHS
+ : I32EnumAttrCase<"MATMUL_F16F16F32_LHS", 6>;
+def MATMUL_F16F16F32_RHS
+ : I32EnumAttrCase<"MATMUL_F16F16F32_RHS", 7>;
+def MATMUL_F16F16F32_RESULT
+ : I32EnumAttrCase<"MATMUL_F16F16F32_RESULT", 8>;
+def MATMUL_F16F16F16_LHS
+ : I32EnumAttrCase<"MATMUL_F16F16F16_LHS", 9>;
+def MATMUL_F16F16F16_RHS
+ : I32EnumAttrCase<"MATMUL_F16F16F16_RHS", 10>;
+def MATMUL_F16F16F16_RESULT
+ : I32EnumAttrCase<"MATMUL_F16F16F16_RESULT", 11>;
+def MATMUL_BF16BF16F32_LHS
+ : I32EnumAttrCase<"MATMUL_BF16BF16F32_LHS", 12>;
+def MATMUL_BF16BF16F32_RHS
+ : I32EnumAttrCase<"MATMUL_BF16BF16F32_RHS", 13>;
+def MATMUL_BF16BF16F32_RESULT
+ : I32EnumAttrCase<"MATMUL_BF16BF16F32_RESULT", 14>;
+def MATMUL_BF16BF16BF16_LHS
+ : I32EnumAttrCase<"MATMUL_BF16BF16BF16_LHS", 15>;
+def MATMUL_BF16BF16BF16_RHS
+ : I32EnumAttrCase<"MATMUL_BF16BF16BF16_RHS", 16>;
+def MATMUL_BF16BF16BF16_RESULT
+ : I32EnumAttrCase<"MATMUL_BF16BF16BF16_RESULT", 17>;
def TensorEncodingEnum
: I32EnumAttr<"TensorEncoding",
"identifier for encoding used for the tensor",[
MATMUL_F32F32F32_LHS, MATMUL_F32F32F32_RHS, MATMUL_F32F32F32_RESULT,
MATMUL_I8I8I32_LHS, MATMUL_I8I8I32_RHS, MATMUL_I8I8I32_RESULT,
+ MATMUL_F16F16F32_LHS, MATMUL_F16F16F32_RHS, MATMUL_F16F16F32_RESULT,
+ MATMUL_F16F16F16_LHS, MATMUL_F16F16F16_RHS, MATMUL_F16F16F16_RESULT,
+ MATMUL_BF16BF16F32_LHS, MATMUL_BF16BF16F32_RHS, MATMUL_BF16BF16F32_RESULT,
+ MATMUL_BF16BF16BF16_LHS, MATMUL_BF16BF16BF16_RHS, MATMUL_BF16BF16BF16_RESULT,
]> {
let cppNamespace = "::mlir::iree_compiler::IREE::LinalgExt";
let genSpecializedAttr = 0;
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/MaterializeEncoding.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/MaterializeEncoding.cpp
index 86ec963..53db781 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/MaterializeEncoding.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/MaterializeEncoding.cpp
@@ -203,15 +203,7 @@
getEncoding(inputs[1]->get().getType().cast<RankedTensorType>());
std::optional<TensorEncoding> resultEncoding =
getEncoding(outputs[0]->get().getType().cast<RankedTensorType>());
- if (!lhsEncoding ||
- (lhsEncoding.value() != TensorEncoding::MATMUL_F32F32F32_LHS &&
- lhsEncoding.value() != TensorEncoding::MATMUL_I8I8I32_LHS) ||
- !rhsEncoding ||
- (rhsEncoding.value() != TensorEncoding::MATMUL_F32F32F32_RHS &&
- rhsEncoding.value() != TensorEncoding::MATMUL_I8I8I32_RHS) ||
- !resultEncoding ||
- (resultEncoding.value() != TensorEncoding::MATMUL_F32F32F32_RESULT &&
- resultEncoding.value() != TensorEncoding::MATMUL_I8I8I32_RESULT)) {
+ if (!lhsEncoding || !rhsEncoding || !resultEncoding) {
return failure();
}
Operation *mmt4DOp = rewriter.create<linalg::Mmt4DOp>(
