add int support and fix lint

Author: Siyuan Feng

src/codegen/codegen_cuda.cc

@@ -389,11 +389,11 @@ void CodeGenCUDA::VisitStmt_(const Allocate* op) {
     std::string scope = alloc_storage_scope_.at(buffer);
     if (scope.find("wmma.") == 0) {
       if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") {
-        CHECK(op->type.is_float() && op->type.bits() == 16)
-          << "Matrix_a and matrix_b only support half type for now";
+        CHECK(op->type == Float(16) || op->type == Int(8) || op->type == UInt(8))
+          << "Matrix_a and matrix_b only support half or char or unsigned char type for now";
       } else {
-        CHECK(op->type.is_float() && (op->type.bits() == 16 || op->type.bits() == 32))
-          << "Accumulator only support half and float type for now";
+        CHECK(op->type == Float(16) || op->type == Float(32) || op->type == Int(32))
+          << "Accumulator only support half, float and int type for now";
       }
       constant_size /= 256;
       PrintWmmaScope(scope, op->type, buffer, stream);
@@ -511,7 +511,8 @@ void CodeGenCUDA::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*
   PrintConst(op, os, this);
 }
 
-void CodeGenCUDA::PrintWmmaScope(const std::string &scope, Type t, const Variable* variable, std::ostream &os) {
+void CodeGenCUDA::PrintWmmaScope(const std::string &scope, Type t,
+    const Variable* variable, std::ostream &os) {
   std::stringstream type;
   PrintType(t, type);
   std::string shape_str = fragment_shapes[variable];
@@ -527,7 +528,8 @@ void CodeGenCUDA::PrintWmmaScope(const std::string &scope, Type t, const Variabl
        << shape_str << ", " << type.str() << ", nvcuda::wmma::" << layout_str <<">";
   } else if (scope == "wmma.accumulator") {
     need_mma_h_ = true;
-    os << "nvcuda::wmma::fragment<nvcuda::wmma::accumulator, " << shape_str << ", "<< type.str() << ">";
+    os << "nvcuda::wmma::fragment<nvcuda::wmma::accumulator, "
+       << shape_str << ", "<< type.str() << ">";
   }
 }
 

src/codegen/codegen_cuda.h

@@ -28,6 +28,7 @@
 #include <tvm/codegen.h>
 #include <tvm/packed_func_ext.h>
 #include <string>
+#include <unordered_map>
 #include "codegen_c.h"
 
 namespace tvm {

src/pass/infer_fragment.cc

@@ -119,7 +119,7 @@ class FragmentGetter : public IRVisitor {
 
 class FragmentChecker : public IRVisitor {
  public:
-  FragmentChecker(const FragmentGetter &getter) : fragment_getter(getter) {}
+  explicit FragmentChecker(const FragmentGetter &getter) : fragment_getter(getter) {}
 
   void Visit_(const Call* op) final {
     if (op->is_intrinsic(intrinsic::tvm_mma_sync)) {
@@ -137,22 +137,22 @@ class FragmentChecker : public IRVisitor {
       CHECK(CheckShape(buffer_var_d, buffer_var_c));
     }
   }
+
  private:
   bool CheckShape(const Variable* buffer1, const Variable* buffer2) {
     CHECK(fragment_getter.fragments.count(buffer1));
     CHECK(fragment_getter.fragments.count(buffer2));
     FragmentGetter::FragmentInfo info1 = fragment_getter.fragments.at(buffer1);
     FragmentGetter::FragmentInfo info2 = fragment_getter.fragments.at(buffer2);
     return info1.m == info2.m && info1.n == info2.n && info1.k == info2.k;
-
   }
-  const FragmentGetter &fragment_getter;
 
+  const FragmentGetter &fragment_getter;
 };
 
 class InferFragmenter : public IRMutator {
  public:
-  InferFragmenter(const FragmentGetter &getter) : fragment_getter(getter) {}
+  explicit InferFragmenter(const FragmentGetter &getter) : fragment_getter(getter) {}
 
   Stmt Mutate_(const Allocate* op, const Stmt& s) final {
     Stmt stmt = IRMutator::Mutate_(op, s);
@@ -174,6 +174,7 @@ class InferFragmenter : public IRMutator {
     }
     return stmt;
   }
+
  private:
   const FragmentGetter &fragment_getter;
 };