[TPP-Pipeline] New flag option to disable vnni packing for packed types. (#1077)

This `patch` adds new flag option to disable vnni layout for packed types.

Author: arun-thmn

include/TPP/PassBundles.td

@@ -52,6 +52,9 @@ def DefaultTppPasses : Pass<"default-tpp-passes", "ModuleOp"> {
     Option<"lowerPackUnpackWithoutTranspose", "lower-pack-unpack-without-transpose",
            "bool", /*default=*/"false",
            "Lower non-constant packs and unpacks reverting any dim permutations.">,
+    Option<"disableVnniPacking", "disable-vnni-packing",
+           "bool", /*default=*/"false",
+           "Disables VNNI packing for packed types.">,
     ListOption<"registerBlocking", "registerBlocking",
            "unsigned", "Register blocking tile sizes for brgemm operation.">,
 
@@ -71,7 +74,10 @@ def TppMapping : Pass<"tpp-mapping", "ModuleOp"> {
   let options= [
     Option<"lowerPackUnpackWithoutTranspose", "lower-pack-unpack-without-transpose",
            "bool", /*default=*/"false",
-           "Lower non-constant packs and unpacks reverting any dim permutations.">
+           "Lower non-constant packs and unpacks reverting any dim permutations.">,
+    Option<"disableVnniPacking", "disable-vnni-packing",
+           "bool", /*default=*/"false",
+           "Disables VNNI packing for packed types.">
   ];
 }
 

lib/TPP/DefaultPipeline.cpp

@@ -67,6 +67,9 @@ llvm::cl::opt<bool> lowerPackUnpackWithoutTranspose(
     llvm::cl::desc("Lower packs and unpacks reverting any dim permutations"),
     llvm::cl::init(false));
 
+llvm::cl::opt<bool> disableVnniPacking("disable-vnni-packing",
+                                   llvm::cl::desc("Disables VNNI packing for packed types"),
+                                   llvm::cl::init(false));
 
 llvm::cl::list<unsigned>
     registerBlocking("registerBlocking", llvm::cl::desc("Register blocking tile sizes for brgemm operation"),
@@ -156,6 +159,7 @@ struct DefaultPipeline : public tpp::impl::DefaultPipelineBase<DefaultPipeline>,
       tppDefaultOptions.linalgToVector = linalgToVector;
       tppDefaultOptions.vectorToXSMM = vectorToXSMM;
       tppDefaultOptions.lowerPackUnpackWithoutTranspose = lowerPackUnpackWithoutTranspose;
+      tppDefaultOptions.disableVnniPacking = disableVnniPacking;
       tppDefaultOptions.registerBlocking =
           SmallVector<unsigned>{registerBlocking.begin(), registerBlocking.end()};
       tppDefaultOptions.vectorToKernel = vectorToKernel;

lib/TPP/DefaultTppPasses.cpp

@@ -117,7 +117,8 @@ struct DefaultTppPasses
       pm.addPass(createRewriteBatchMatmulToMatmul());
 
       // Applies a set of passes at the linalg level to fuse and pack.
-      TppMappingOptions tppMappingOptions{lowerPackUnpackWithoutTranspose};
+      TppMappingOptions tppMappingOptions{lowerPackUnpackWithoutTranspose, 
+	        disableVnniPacking};
       pm.addPass(createTppMapping(tppMappingOptions));
 
       // Generalize linalg.pack and linalg.unpack.

lib/TPP/PassBundles/TppMapping.cpp

@@ -62,7 +62,10 @@ struct TppMapping : public tpp::impl::TppMappingBase<TppMapping>,
     pm.addPass(createPackConv2DNchwFchw());
     pm.addPass(createRewriteConvToMatmulOrBrgemm());
     pm.addPass(createPackMatmul());
-    pm.addPass(createPackVNNI());
+
+    if (!disableVnniPacking) {
+      pm.addPass(createPackVNNI());
+    }
 
     if (lowerPackUnpackWithoutTranspose) {
       pm.addPass(createLowerPacksAndUnpacksWithoutTranspose());