[quantization][opencl]: Support quantization for SplatInst

pytorch · Jul 23, 2018 · ef071c9 · ef071c9
1 parent 04fb697
commit ef071c9
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/lib/Backends/OpenCL/OpenCL.cpp b/lib/Backends/OpenCL/OpenCL.cpp
@@ -566,7 +566,16 @@ void OpenCLFunction::execute() {
 
       if (auto *SI = dyn_cast<SplatInst>(&I)) {
         // Pass the splat as a parameter.
-        setKernelArg(kernel, ++numArgs, SI->getValue());
+        if (!isQuantized) {
+          setKernelArg(kernel, ++numArgs, SI->getValue());
+        } else {
+          auto *destTy = SI->getDest()->getType();
+          TensorQuantizationParams destQ{destTy->getScale(),
+                                         destTy->getOffset()};
+          float val = SI->getValue();
+          int8_t int8Val = quantization::quantize(val, destQ);
+          setKernelArg<int8_t>(kernel, ++numArgs, int8Val);
+        }
       } else if (auto *EPI = dyn_cast<ElementPowInst>(&I)) {
         // Pass the exp as a parameter.
         setKernelArg(kernel, ++numArgs, EPI->getExp());

diff --git a/lib/Backends/OpenCL/OpenCL.h b/lib/Backends/OpenCL/OpenCL.h
@@ -186,6 +186,7 @@ class OCLBackend final : public Backend {
       case Kinded::Kind::MinNodeKind:
       case Kinded::Kind::MulNodeKind:
       case Kinded::Kind::QuantizeNodeKind:
+      case Kinded::Kind::SplatNodeKind:
       case Kinded::Kind::SubNodeKind:
       case Kinded::Kind::TransposeNodeKind:
         return true;

diff --git a/lib/Backends/OpenCL/kernels.cl b/lib/Backends/OpenCL/kernels.cl
@@ -467,6 +467,7 @@ DEFINE_OPENCL_TERNARY_DATA_PARALLEL_KERNEL(elementselect, float,
 
 DEFINE_OPENCL_UNARY_DATA_PARALLEL_KERNEL_WITH_IMM_OPERAND(splat, float, SRC)
 DEFINE_OPENCL_UNARY_DATA_PARALLEL_KERNEL_WITH_IMM_OPERAND(splat_u, ulong, SRC)
+DEFINE_OPENCL_UNARY_DATA_PARALLEL_KERNEL_WITH_IMM_OPERAND(splat_i8, char, SRC)
 
 #undef DEFINE_OPENCL_BINARY_DATA_PARALLEL_KERNEL_WITH_IMM_OPERAND
 #undef DEFINE_OPENCL_UNARY_DATA_PARALLEL_KERNEL_WITH_IMM_OPERAND