mlir/test/Conversion/ArithToAMDGPU/8-bit-float-saturation.mlir

   1 // RUN: mlir-opt --split-input-file %s \
   2 // RUN:   --pass-pipeline='builtin.module(func.func(convert-arith-to-amdgpu{chipset=gfx940 saturate-fp8-truncf=true}))' \
   3 // RUN:   | FileCheck %s
   4
   5 // CHECK-LABEL: func.func @scalar_trunc
   6 // CHECK-SAME: ([[V:%.+]]: f16)
   7 // CHECK-DAG: [[CMin:%.+]] = arith.constant -5.734400e+04 : f16
   8 // CHECK-DAG: [[CMax:%.+]] = arith.constant 5.734400e+04 : f16
   9 // CHECK-DAG: [[CInf:%.+]] = arith.constant 0x7C00 : f16
  10 // CHECK-DAG: [[CNegInf:%.+]] = arith.constant 0xFC00 : f16
  11 // CHECK: [[ISINF:%.+]] = arith.cmpf oeq, [[V]], [[CInf]]
  12 // CHECK: [[ISNEGINF:%.+]] = arith.cmpf oeq, [[V]], [[CNegInf]]
  13 // CHECK: [[ISNAN:%.+]] = arith.cmpf uno, [[V]], [[V]]
  14 // CHECK: [[ISNONFINITE_1:%.+]] = arith.ori [[ISINF]], [[ISNEGINF]]
  15 // CHECK: [[ISNONFINITE:%.+]] = arith.ori [[ISNONFINITE_1]], [[ISNAN]]
  16 // CHECK: [[CLAMPEDBELOW:%.+]] = arith.maximumf [[V]], [[CMin]]
  17 // CHECK: [[CLAMPED:%.+]] = arith.minimumf [[CLAMPEDBELOW]], [[CMax]]
  18 // CHECK: [[SATURATED:%.+]] = arith.select [[ISNONFINITE]], [[V]], [[CLAMPED]]
  19 // CHECK: [[FLOAT:%.+]] = arith.extf [[SATURATED]] : f16 to f32
  20 // CHECK: [[TRUNCV:%.+]] = amdgpu.packed_trunc_2xfp8 [[FLOAT]], undef into undef[word 0] : f32 to vector<4xf8E5M2FNUZ>
  21 // CHECK: [[W:%.+]] = vector.extract [[TRUNCV]][0] : f8E5M2FNUZ from vector<4xf8E5M2FNUZ>
  22 // CHECK: return [[W]] : f8E5M2FNUZ
  23 func.func @scalar_trunc(%v: f16) -> f8E5M2FNUZ {
  24   %w = arith.truncf %v : f16 to f8E5M2FNUZ
  25   return %w : f8E5M2FNUZ
  26 }
  27
  28 // No 0-D test because arith.truncf hasn't been extended to support it.
  29
  30 // -----
  31
  32 // CHECK-LABEL: func.func @vector_trunc
  33 // CHECK-SAME: ([[V:%.+]]: vector<2xf32>) -> vector<2xf8E4M3FNUZ> {
  34 // CHECK-DAG: [[CMin:%.+]] = arith.constant dense<-2.400000e+02> : vector<2xf32>
  35 // CHECK-DAG: [[CMax:%.+]] = arith.constant dense<2.400000e+02> : vector<2xf32>
  36 // CHECK-DAG: [[CInf:%.+]] = arith.constant dense<0x7F800000> : vector<2xf32>
  37 // CHECK-DAG: [[CNegInf:%.+]] = arith.constant dense<0xFF800000> : vector<2xf32>
  38 // CHECK: [[ISINF:%.+]] = arith.cmpf oeq, [[V]], [[CInf]]
  39 // CHECK: [[ISNEGINF:%.+]] = arith.cmpf oeq, [[V]], [[CNegInf]]
  40 // CHECK: [[ISNAN:%.+]] = arith.cmpf uno, [[V]], [[V]]
  41 // CHECK: [[ISNONFINITE_1:%.+]] = arith.ori [[ISINF]], [[ISNEGINF]]
  42 // CHECK: [[ISNONFINITE:%.+]] = arith.ori [[ISNONFINITE_1]], [[ISNAN]]
  43 // CHECK: [[CLAMPEDBELOW:%.+]] = arith.maximumf [[V]], [[CMin]]
  44 // CHECK: [[CLAMPED:%.+]] = arith.minimumf [[CLAMPEDBELOW]], [[CMax]]
  45 // CHECK: [[SATURATED:%.+]] = arith.select [[ISNONFINITE]], [[V]], [[CLAMPED]]
  46 // CHECK: [[F0:%.+]] = vector.extract [[SATURATED]][0]
  47 // CHECK: [[F1:%.+]] = vector.extract [[SATURATED]][1]
  48 // CHECK: [[W0:%.+]] = amdgpu.packed_trunc_2xfp8 [[F0]], [[F1]] into undef[word 0] : f32 to vector<4xf8E4M3FNUZ>
  49 // CHECK: [[W:%.+]] = vector.extract_strided_slice [[W0]] {offsets = [0], sizes = [2], strides = [1]} : vector<4xf8E4M3FNUZ> to vector<2xf8E4M3FNUZ>
  50 // CHECK: return [[W]] : vector<2xf8E4M3FNUZ>
  51 func.func @vector_trunc_short(%v: vector<2xf32>) -> vector<2xf8E4M3FNUZ> {
  52   %w = arith.truncf %v : vector<2xf32> to vector<2xf8E4M3FNUZ>
  53   return %w : vector<2xf8E4M3FNUZ>
  54 }